题解 | #最小的K个数#

描述
给定一个数组，找出其中最小的K个数。例如数组元素是4,5,1,6,2,7,3,8这8个数字，则最小的4个数字是1,2,3,4。
0 <= k <= input.length <= 10000
0 <= input[i] <= 10000

示例1
输入：
[4,5,1,6,2,7,3,8],4
返回值：
[1,2,3,4]
说明：
返回最小的4个数即可，返回[1,3,2,4]也可以

示例2
输入：
[1],0
返回值：
[]
示例3
输入：
[0,1,2,1,2],3
返回值：
[0,1,1]

class Solution {
public:
    vector<int> GetLeastNumbers_Solution(vector<int> input, int k) {
        //return minK_heap(input, k);      // 堆实现
        return minK_quickSelect(input, k); // 快排实现
    }

    vector<int> minK_heap(vector<int> input, int k) {
        vector<int> ans;
        if (k == 0) return ans;
        priority_queue<int> pq;
        int n = input.size();
        for (int i = 0; i < min(n, k); ++i) {
            pq.push(input[i]);  // 创建大小为 k 的大顶堆
        }
        for (int i = k; i < n; ++i) {
            if (input[i] < pq.top()) { // 新数据小于对顶数据（最大的那个）
                pq.pop();
                pq.push(input[i]); // 更新堆
            }
        }
        ans.reserve(k);  // 分配空间
        while (!pq.empty()) {
            ans.push_back(pq.top());
            pq.pop();
        }
        return ans;
    }

    vector<int> minK_quickSelect(vector<int>& input, int k) {
        vector<int> ret;
        if (k==0 || k > input.size()) return ret;
        int l = 0, r = input.size();
        while (l < r) {
            int p = partition(input, l, r);  // 分段
            if (p+1 == k) { // 找到第 k 大的数据
                return vector<int>({input.begin(), input.begin()+k}); // 返回 k 个最小值
            }
            if (p+1 < k) { // 右半部分
                l = p + 1; 
            }  
            else {         // 左半部分
                r = p;
            }
        }
        return ret;
    }

    int partition(vector<int>& input, int l, int r) {  // 分段实现
        int pivot = input[r-1];
        int i = l;
        for (int j=l; j<r-1; ++j) {
            if (input[j] < pivot) {
                swap(input[i++], input[j]); 
            }
        }
        swap(input[i], input[r-1]);
        return i;
    }
};