算法实战：快速选择解决第 K 大元素与最小 K 个数问题

流程解析

46. 最小的 K 个数

题目描述

输入整数数组 stock 和整数 cnt，找出数组中最小的 cnt 个数。

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

核心思路

这道题同样可以用快速选择算法解决。相比堆排序（O(NlogK)），快速选择的平均复杂度更接近 O(N)。

逻辑上，我们依然采用三路划分。不同的是，我们要找的是最小的 K 个数，所以主要关注**左侧区域（小于基准）**的长度：

1. 如果左侧区域长度 >= cnt，说明目标都在左边，递归处理左半部分。
2. 如果左侧区域长度 < cnt，但左侧 + 中间 >= cnt，说明基准及其左侧已经覆盖了所有需要的数，直接结束。
3. 否则，需要在右侧继续寻找剩余的数。

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));
        }
    }
};

流程解析

总结

这两道题的核心都在于快速选择算法。通过随机选取基准值并进行三路划分，我们可以在不完成全排序的情况下，根据区间长度直接定位目标范围。

相比于传统的全排序（O(NlogN)）或堆排序（O(NlogK)），这种方法在平均情况下能将时间复杂度优化至 O(N)。虽然代码逻辑稍显复杂，需要仔细处理边界条件和递归终止条件，但在处理大规模数据查找时，性能优势非常明显。希望大家能掌握这种分治优化的思路。