剑指offer：最小的k个数

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

示例 1：
输入：arr = [3,2,1], k = 2
输出：[1,2] 或者 [2,1]
示例 2：
输入：arr = [0,1,2,1], k = 1
输出：[0]

\quad 这里给出一个三路partition的操作，速度最快，可见测试数据中有很多重复元素。

class Solution {
public:
    int partition(vector<int> arr, int k){
        int n = arr.size();
        if(n == 1 && k == 1) {
            return arr[0];
        }
        int f = arr[rand() % n]; // 随机取一个标志位
        vector<int> a, b, c;  // 三路partition
        for(auto num: arr){
            if(num < f) a.push_back(num);
            else if(num == f) b.push_back(num);
            else c.push_back(num);
        }
        if(a.size() >= k) return partition(a, k);
        else if(a.size() + b.size() >= k) return f; // 此时第k个数就是答案
        else return partition(c, k - a.size() - b.size());
    }
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        vector<int> res;
        if(arr.empty() || k == 0) return res;
        int f = partition(arr, k);
        for(int num: arr) if(num < f) res.push_back(num);
        while(res.size() < k) res.push_back(f);
        return res;
    }
};