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

流程解析

算法执行时，首先随机选取一个基准值，将数组划分为大于、等于、小于三部分。随后根据各部分的大小关系，决定是继续向右递归、向左递归还是直接返回基准值。这种剪枝策略使得我们不需要对整个数组排序即可找到目标。

46. 最小的 K 个数

题目描述

输入整数数组 stock 和整数 cnt，请返回数组中最小的 cnt 个数。

解法思路

这道题与上一题逻辑高度相似，同样采用快速选择的分区策略。不同之处在于，我们需要保留左侧较小的 cnt 个数。当分区完成后，如果左侧区域（小于基准值的区域）大小已经满足 cnt，则无需再处理右侧；如果不足，则需要在右侧继续寻找剩余数量的元素。

相比堆排序（O(NlogK)）和全排序（O(NlogN)），快速选择算法的平均时间复杂度接近 O(N)，在处理大规模数据时优势明显。

代码实现

class Solution {
public:
    vector<int> inventoryManagement(vector<int>& stock, int cnt) {
        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++;
            }
        }
        
        // 左侧区域元素个数 >= cnt，说明最小的 k 个数都在左侧
        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));
        }
    }
};

流程解析

通过同样的三路划分逻辑，我们不断缩小搜索范围。当左侧累积的元素数量达到 cnt 时，数组的前 cnt 个位置即为所求的最小值集合。这种方法避免了不必要的排序开销，是解决此类 Top-K 问题的经典方案。