《剑指offer》刷题系列——（五十八）最小的k个数

题目

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

思路

一、利用快速排序的思想。
快速排序的partition函数每次返回一个索引值index，一次排序之后，index左边的数都小于index位置上的数。因此当我们返回index值时，就将它与k-1值比较，如果刚好等于k-1，则直接输出arr数组的前k个元素（索引为0到k-1）；如果index小于k-1，此时index左边的值还不够k个，需要对index右边的数再次进行排序；如果index大于k-1，那么我们需要的k个数就是index左边的一部分数，但是还不知道前k个，所以它们还需要排序。

二、最大堆
python中默认是最小堆，所以要对数组中所有的数取其相反数，才能使用小根堆维护前 k 小值。

代码

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        if not arr or k<=0:
            return []
        start = 0
        end = len(arr)-1
        index = self.partition(arr,start,end)
        while index!=k-1:
            if index>k-1:
                end = index-1
                index = self.partition(arr,start,end)
            else:
                start = index+1
                index = self.partition(arr,start,end)

        return arr[:k]
        
    
    def partition(self,arr,left,right):
        key = arr[left]
        while left<right:
            while left<right and arr[right]>=key:
                right-=1
            arr[left]=arr[right]
            while left<right and arr[left]<=key:
                left+=1
            arr[right]=arr[left]
        arr[left]=key
        return left

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        if k == 0:
            return []
            
        hp = [-i for i in arr[:k]]
        heapq.heapify(hp)
        for i in range(k,len(arr)):
            if arr[i]<-hp[0]:
                heapq.heappop(hp)
                heapq.heappush(hp,-arr[i])
        res = [-x for x in hp]
        return res

复杂度

一、快速排序
平均时间复杂度为 O(N)
因为我们是要找下标为k的元素，第一次切分的时候需要遍历整个数组 (0 ~ n) 找到了下标是 index的元素，假如 k 比 index小的话，那么我们下次切分只要遍历数组 (0~k-1)的元素就行啦，反之如果 k 比 j 大的话，那下次切分只要遍历数组 (k+1～n) 的元素就行啦，总之可以看作每次调用 partition 遍历的元素数目都是上一次遍历的 1/2，因此时间复杂度是 N + N/2 + N/4 + … + 1 = 2N, 因此时间复杂度是 O(N)。

空间复杂度：期望为 O(log⁡n)，递归调用的期望深度为 O(log⁡n)，每层需要的空间为 O(1)

二、最大堆
时间复杂度：O(nlog⁡k)，其中 n 是数组 arr 的长度。由于大根堆实时维护前 k 小值，所以插入删除都是 O(log⁡k)的时间复杂度，最坏情况下数组里 n 个数都会插入，所以一共需要 O(nlog⁡k) 的时间复杂度。

空间复杂度：O(k)，因为大根堆里最多 k个数。

知识点：关于python中堆的使用（默认为最小堆）

在堆中插入元素

>>> import heapq
>>> h = []
>>> heapq.heappush(h,2)
>>> h
[2]

将列表转换为堆

>>> list = [1,2,3,5,4,3,1,8,6,9,0]
>>> heapq.heapify(list)
>>> list
[0, 1, 1, 5, 2, 3, 3, 8, 6, 9, 4]
>>>

删除并返回最小值，因为堆的特征是heap[0]永远是最小的元素，所以一般都是删除第一个元素

>>> list
[0, 1, 1, 5, 2, 3, 3, 8, 6, 9, 4]
>>> heapq.heappop(list)
0
>>> list
[1, 1, 3, 5, 2, 3, 4, 8, 6, 9]
>>>

删除并返回最小元素值，添加新的元素值

>>> list
[1, 1, 3, 5, 2, 3, 4, 8, 6, 9]
>>> heapq.heapreplace(list,88)
1
>>> list
[1, 2, 3, 5, 9, 3, 4, 8, 6, 88]
>>>

查询堆中的最大n个元素或最小n个元素

>>> list
[1, 2, 3, 5, 9, 3, 4, 8, 6, 88]
>>> heapq.nlargest(4,list)
[88, 9, 8, 6]
>>> heapq.nsmallest(3,list)
[1, 2, 3]
>>>

判断添加元素值item与堆的第一个元素值对比，如果大，则删除并返回第一个元素，然后添加新元素值item。 如果小，则返回item。原堆不变。

>>> list
[1, 2, 3, 5, 9, 3, 4, 8, 6, 88]
>>> heapq.heappushpop(list,6)
1
>>> list
[2, 5, 3, 6, 9, 3, 4, 8, 6, 88]
>>> heapq.heappushpop(list,0)
0
>>> list
[2, 5, 3, 6, 9, 3, 4, 8, 6, 88]