【剑指Offer系列40】最小的k个数

题目

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

代码

Python

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        if k==0: return list()
        
        hp = [-x for x in arr[:k]] # 准备建立最大堆，负号因为后面调用堆的方法弹出最小值
        heapq.heapify(hp) # 将列表转换为堆
        # 遍历列表剩余元素，最后堆中留下的就是那k个最小值
        for i in range(k,len(arr)):
            if -hp[0] > arr[i]: # 堆顶元素大于待插入元素
                heapq.heappop(hp) # 弹出堆内最小值
                heapq.heappush(hp,-arr[i]) # 压入待插入元素的相反数
        res = [-x for x in hp]
        return res

C++

class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        vector<int> vec(k,0); // 创建一个有k个0的数组
        int n=arr.size();
        if (k==0) return vec;
        
        priority_queue<int> hp; // 堆就是一种优先队列
        for (int i=0; i<k; ++i) hp.push(arr[i]); // 先建立k个数的堆,采用++i比i++更快
        for (int j=k; j<n; ++j) {
            if (hp.top() > arr[j]) {
                hp.pop();
                hp.push(arr[j]);
            }
        }
        for (int m=0; m<k; ++m) {
            vec[m]=hp.top();
            hp.pop();
        }
        return vec;
    }
};