算法实战：快速选择解决第 K 大元素与最小 K 个数问题
在处理数组查找类问题时，全排序往往显得过于重手。今天我们来聊聊如何利用**快速选择算法（Quick Select）**高效解决两个经典场景：数组中的第 K 个最大元素，以及最小的 K 个数。
45. 数组中的第 K 个最大元素
题目描述
给定一个整数数组 nums 和一个整数 k，返回该数组中第 k 个最大的元素。
核心思路
很多人第一反应是排序后取索引，但时间复杂度是 O(NlogN)。其实我们可以利用快排的分区思想进行优化。
在标准的三路快排中，我们将数组划分为 [小于基准]、[等于基准]、[大于基准] 三个区域。对于找第 K 大的数，我们只需要关注这三个区域的长度：
- 如果**右侧区域（大于基准）**的元素个数 >= k，说明目标就在右边，递归进入右区间。
- 如果**中间区域（等于基准）**包含了目标位置，那么基准值本身就是答案。
- 否则，目标在左侧，调整 k 值后递归左区间。
这种策略下，平均时间复杂度可以逼近 O(N)。
C++ 代码实现
class Solution {
public:
int Top_k(vector<int>& nums, int left, int right, int k) {
if (left == right) {
return nums[left];
}
int l = left - 1, r = right + 1, i = left;
// 随机选择基准元素，避免最坏情况
int key = nums[rand() % (right - left + 1) + left];
// 三路划分：大于 key | 等于 key | 小于 key
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 个数。
核心思路
这道题同样可以用快速选择算法解决。相比堆排序（O(NlogK)），快速选择的平均复杂度更接近 O(N)。
逻辑上，我们依然采用三路划分。不同的是，我们要找的是最小的 K 个数，所以主要关注**左侧区域（小于基准）**的长度：
- 如果左侧区域长度 >= cnt，说明目标都在左边，递归处理左半部分。
- 如果左侧区域长度 < cnt，但左侧 + 中间 >= cnt，说明基准及其左侧已经覆盖了所有需要的数，直接结束。
- 否则，需要在右侧继续寻找剩余的数。
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)。虽然代码逻辑稍显复杂，需要仔细处理边界条件和递归终止条件，但在处理大规模数据查找时，性能优势非常明显。希望大家能掌握这种分治优化的思路。