【剑指Offer】面试题40. 最小的k个数（c++ 优先队列 [priority_queue] / 第n大元素线性时间选择 [nth_element]）

题目

解法一：选择排序 (前n)

O(n * k)

class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        vector<int> ans(k);

        for(int i = 0; i < k; i++){
            int min = 0;
            for(int j = 1; j < arr.size(); j++){
                if(arr[j] < arr[min]){
                    min = j;
                }
            }
            ans[i] = arr[min];
            arr[min] = INT_MAX;
        }

        return ans;
    }
};

解法二：大根堆

O(n * log k)

class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        vector<int> ans(k);
        priority_queue<int> heap; // c++ 里这东西是用大顶堆维护的

        if(k == 0) return ans;

        // 把前k个放进去
        for(int i = 0; i < k; i++){
            heap.push(arr[i]);
        }

        // 当新元素堆顶大，则移除堆顶并插入，堆始终维护当前前k小的值
        for(int i = k; i < arr.size(); i++){
            if(arr[i] < heap.top()){
                heap.pop();
                heap.push(arr[i]);
            }
        }

        // 放到数组中
        for(int i = 0; i < k; i++){
            ans[i] = heap.top();
            heap.pop();
        }

        
        return ans;
    }
};

解法三：线性时间选择（使用快排的选择函数）

O(n)

class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        vector<int> ans(k);

        // 判断 0
        if(k == 0){
            return ans;
        }

        //  std::nth_element(std::begin(numbers), std::begin(numbers) + count, std::end(numbers));
        //  该函数属于algorithm库
        //  参数1：容器迭代器begin
        //  参数2：容器迭代器begin + count
        //  参数3：容器迭代器end
        //  功能：将count位置元素归位，左侧均小于该元素，右侧均大于该元素
        //  原理：利用快速选择的线性时间选择，过程与快速排序类似，只不过只进入其中一个分支
        nth_element(arr.begin(), arr.begin() + k, arr.end());
        
        // 前k个元素即为答案
        for(int i = 0; i < k; i++){
            ans[i] = arr[i];
        }

        return ans;
    }
};