avatar

Catalog
面试题40. 最小的k个数

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

我们用一个大根堆实时维护数组的前 kk 小值。首先将前 kk 个数插入大根堆中,随后从第 k+1k+1 个数开始遍历,如果当前遍历到的数比大根堆的堆顶的数要小,就把堆顶的数弹出,再插入当前遍历到的数。最后将大根堆里的数存入数组返回即可。在下面的代码中,由于 C++ 语言中的堆(即优先队列)为大根堆,我们可以这么做。而 Python 语言中的对为小根堆,因此我们要对数组中所有的数取其相反数,才能使用小根堆维护前 kk 小值。

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        if k == 0:
            return list()
        heap = [-arr[i] for i in range(k)]
        heapq.heapify(heap)
        for i in range(k, len(arr)):
            if -heap[0] > arr[i]:
                heapq.heappop(heap)
                heapq.heappush(heap, -arr[i])
        ans = [-x for x in heap]
        return ans       
Author: kim yhow
Link: http://yoursite.com/2020/03/20/面试题40-最小的k个数/
Copyright Notice: All articles in this blog are licensed under CC BY-NC-SA 4.0 unless stating additionally.
Donate
  • 微信
    微信
  • 支付寶
    支付寶