快速选择算法实战：数组中第 K 个最大元素与最小的 K 个数

算法总结及流程解析

46. 最小的 k 个数

题目链接

LCR 159. 库存管理 III - 力扣（LeetCode）

题目描述

题目示例

解法 (快速选择算法)

算法思路

同样利用快排思想，当我们把数组「分成三块」之后：[ l, left ] [ left + 1, right -1 ] [ right, r ]，我们可以通过计算每一个区间内元素的「个数」，进而推断出最小的 k 个数在哪些区间里面。

那么我们可以直接去「相应的区间」继续划分数组即可，无需处理剩余部分。

C++ 算法代码

class Solution {
public:
    vector<int> inventoryManagement(vector<int>& stock, int cnt) {
        // 解法一：快排 (优点：简单无脑；缺点：时间复杂度很大 O(NlogN))
        // sort(stock.begin(), stock.end());
        // vector<int> ret;
        // for(int i = 0; i < cnt; i++) { ret.push_back(stock[i]); }
        // return ret;
        
        // 解法二：堆排序 (优点：时间复杂度比快排小 O(Nlogk)；缺点：比较难想)
        // vector<int> ret;
        // if(cnt == 0) return {};
        // priority_queue<int> pq(stock.begin(), stock.begin() + cnt);
        // for(int i = cnt; i < stock.size(); i++) {
        //     if(pq.top() > stock[i]) { pq.pop(); pq.push(stock[i]); }
        // }
        // while(!pq.empty()) { ret.push_back(pq.top()); pq.pop(); }
        // return ret;
        
        // 解法三：快速选择排序 (优点：时间复杂度非常小：逼近 O(N)；缺点：方法很巧妙很难想到)
        if(cnt == 0) {
            return {};
        }
        srand(time(NULL));
        Top_k(stock, 0, stock.size() - 1, cnt);
        return vector<int>(stock.begin(), stock.begin() + cnt);
    }
    
    void Top_k(vector<int>& nums, int left, int right, int cnt) {
        if(left == right) {
            return;
        }
        int key = nums[rand() % (right - left + 1) + left];
        int l = left - 1, r = right + 1, i = left;
        
        while(i < r) {
            if(nums[i] > key) {
                swap(nums[i], nums[--r]);
            } else if(nums[i] < key) {
                swap(nums[i++], nums[++l]);
            } else {
                i++;
            }
        }
        
        if(l - left + 1 >= cnt) {
            return Top_k(nums, left, l, cnt);
        } else if(r - left >= cnt) {
            return;
        } else {
            return Top_k(nums, r, right, cnt - (r - left));
        }
    }
};

算法总结及流程解析