剑指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;
    }
};
相关推荐
©️2020 CSDN 皮肤主题: 编程工作室 设计师:CSDN官方博客 返回首页