剑指 Offer 40. 最小的k个数(快排思想/大根堆)

题目描述

输入整数数组 arr ，找出其中最小的 k 个数。例如，输入4、5、1、6、2、7、3、8这8个数字，则最小的4个数字是1、2、3、4。

 

示例 1：

输入：arr = [3,2,1], k = 2
输出：[1,2] 或者 [2,1]
示例 2：

输入：arr = [0,1,2,1], k = 1
输出：[0]
 

限制：

0 <= k <= arr.length <= 10000
0 <= arr[i] <= 10000

解法一：快排

这里的快排，是指用快排思想，其实并不需要真正的排序。还记得快排的中值吗，当找到中值后，中值左边的数都比右边的数小，中值右边的数都比左边的数大，那么，我们只需要判断需要找的k个数的k-1的索引与中值的索引的关系：

1.中值索引等于k-1，那么此时中值左边刚好就是要找的k个数；

2.中值索引>k-1，那么说明这k个数在中值的左半部分，我们只需要对中值左边部分继续分界，直到中值索引等于k-1

3.中值索引<k-1，那么说明这个k个数在中值右半部分，我们只需要对中值右边的部分继续分界，直到中值索引等于k-1

那么这里有两种代码实现，一种是对结果排好序的，一种是没排好序的

没排好序的：

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        start = 0
        end = len(arr) - 1
        point = self.quicksort(arr, start, end)

        while point != k-1:
            if point < k -1:
                start = point + 1
                point = self.quicksort(arr, start, end)
            if point > k-1:
                end = point - 1
                point = self.quicksort(arr, start, end)
        return arr[:k]

    def quicksort(self, alist, first, last):
        middlepoint = alist[first]
        leftmark = first + 1
        rightmark = last
        done = False
        while not done:
            while leftmark <= rightmark and alist[leftmark] <= middlepoint:
                leftmark += 1
            while rightmark >= leftmark and alist[rightmark] >= middlepoint:
                rightmark -= 1
            if leftmark > rightmark:
                done = True
            else:
                alist[leftmark], alist[rightmark] = alist[rightmark], alist[leftmark]
        alist[first], alist[rightmark] = alist[rightmark], alist[first]
        return rightmark  # 中值的索引

排好序的：

class Solution(object):
    def getLeastNumbers(self, arr, k):
        """
        :type arr: List[int]
        :type k: int
        :rtype: List[int]
        """
        # 方法一:partition方法(基于快速排序)
        if k > len(arr) or k <= 0:
            return [] 
        start = 0
        end = len(arr) - 1
        index = self.quickSort(arr, start, end)
        while index != k-1:
            # print(index)
            if index > k-1:
                end = index - 1
                index = self.quickSort(arr, start, end)
            if index < k-1:
                start = index + 1
                index = self.quickSort(arr, start, end)
        return arr[:k]

    def quickSort(self, arr, start, end):
        low = start
        high = end
        temp = arr[start]
        while low < high:
            while low < high and arr[high] >= temp:
                high -= 1
            arr[low] = arr[high]
            while low <high and arr[low] < temp:
                low += 1
            arr[high] = arr[low]
        arr[low] = temp
        return low

时间复杂度O(N),最坏情况为O(N^2)

空间复杂度O(log n)

递归调用期望深度为O(log n)

大根堆

大根堆的是一个队列形式实现，堆顶保存的是始终是最大的数，那么我们可以这样，先将数组前k个数放到大根堆里面去，再从k+1依次遍历，假如新的数比堆顶下，那么把堆顶弹出，把这个数存到大根堆里面去，那么一趟遍历下来，大根堆始终维护着最小的k个数。

但是在Python里面没有大根堆，python的实现是小根堆，可以调用heapq。那么呢，需要把这个数组取反存在小根堆去，小根堆堆顶始终维护者最小的数，那么它的相反数其实在真是的数中就是最大的数，那么同理，遍历k+1，假如新的数比小根堆顶元素的相反数小的话，那么就存在小根堆里面去，然后把小根堆的堆顶pop出去。记得return的时候也需要取相反数。

from heapq import *
class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        '''
        python小根堆
        '''
        if k == 0:
            return []
        hp = [-x for x in arr[:k]]
        heapq.heapify(hp)
        for i in range(k, len(arr)):
            if -hp[0] > arr[i]:
                heapq.heappop(hp)
                heapq.heappush(hp, -arr[i])
        return [-x for x in hp]

时间复杂度O(n log k)，其中n是数组长度

空间复杂度O(k)，小根堆里面最多有k个数

剑指 Offer 40. 最小的k个数(快排思想/大根堆)

相关推荐