最小的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]

思路一

堆排序，使用优先队列维护一个大顶堆，每次保证堆的大小是k

代码

class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        int n = arr.size();
        vector<int> ans;
        if(k == 0 || n == 0) return ans;
        priority_queue<int> q;
        for(int i = 0;i < k;i++) q.push(arr[i]);
        for(int i = k;i < n;i++){
            if(arr[i] < q.top()){
                q.pop();
                q.push(arr[i]);
            }
        }
        while(k--){
            ans.push_back(q.top());
            q.pop();
        }
        return ans;
    }
};