在处理'第 K 大'或'最小 K 个数'这类问题时,全排序往往效率不高。这里介绍一种基于快速排序思想优化的快速选择算法(Quick Select),平均时间复杂度可降至 O(N)。
45. 数组中的第 K 个最大元素
题目描述
给定整数数组 nums 和整数 k,请返回该数组中第 k 个最大的元素。
注意你需要找的是排序后第 k 大的元素,而不是第 k 个不同的元素。
解法思路
核心在于利用快速排序的分区(Partition)逻辑。在快排中,我们将数组分为三块:大于基准值、等于基准值、小于基准值。通过计算每个区间的长度,我们可以直接判断目标元素落在哪个区间,从而避免对另一侧进行递归处理。
对于第 K 大元素,我们关注的是右侧(较大值)区域的元素数量。如果右侧区域元素个数 >= k,说明目标在右边;如果在中间区域,则当前基准值即为答案;否则在左边,且需要调整 k 的值。
代码实现
class Solution {
public:
int Top_k(vector<int>& nums, int left, int right, int k) {
if (left == right) {
return nums[left];
}
// 三路划分:[left, l] [l+1, r-1] [r, right]
int l = left - 1, r = right + 1, i = left;
// 随机选择基准元素,避免最坏情况
int key = nums[rand() % (right - left + 1) + 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++;
}
}
// 若右边区域元素个数>=k,说明第 k 大的数在右边区域
if (right - r + 1 >= k) {
return Top_k(nums, r, right, k);
}
// 若右边区域个数<k,但中间加右边区域个数>=k,说明第 k 大的数在中间区域
else if (right - l >= k) {
return key;
}
// 若中间加右边区域个数<k,说明第 k 大的数在左边区域
else {
return Top_k(nums, left, l, k - (right - l));
}
}
int findKthLargest(vector<int>& nums, int k) {
srand(time(NULL));
return Top_k(nums, 0, nums.size() - 1, 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 问题的经典方案。


