剑指Offer系列—— 40.最小的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]

限制

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

思路

我们用一个大根堆维护数组中最小的几个数。首先放进k个数，再看剩余的，若比堆顶元素小，则弹出堆顶元素并将此元素压入，最后此优先队列中就是数组中最小的几个数

代码

class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        if (k == 0) return {};
        priority_queue<int> p;
        for (int i = 0 ;i < k; i ++ ) {
            p.push(arr[i]);
        }
        for (int i = k ;i < arr.size(); i ++ ) {
            if (arr[i] < p.top()) {
                p.pop();
                p.push(arr[i]);
            }
        }
        vector<int> res(k);
        for (int i = 0; i < k; i ++ ) {
            res[i] = p.top();
            p.pop();
        }
        return res;
    }
};