LeetCode-215. 数组中的第K个最大元素
1、题目描述
给定整数数组 nums 和整数 k，请返回数组中第 k 个最大的元素。
请注意，你需要找的是数组排序后的第 k 个最大的元素，而不是第 k 个不同的元素。
你必须设计并实现时间复杂度为 O(n) 的算法解决此问题。
示例 1:
输入: [3,2,1,5,6,4], k = 2
输出: 5
示例 2:
输入: [3,2,3,1,2,4,5,5,6], k = 4
输出: 4
2、代码实现
会超时的代码
class Solution
{
public:
int findKthLargest(vector<int>& nums, int k)
{
int target = k - 1;
int left = 0, right = nums.size() - 1;
while (left <= right) {
int pos = partition(nums, left, right);
if (pos == target) {
return nums[pos];
}
else if (pos < target) {
left = pos + 1;
}
else {
right = pos - 1;
}
}
return -1;
}
int partition(vector<int>& nums, int left, int right)
{
static bool initSrand = false;
if (!initSrand) {
srand(time(0));
}
int pivot_index = left + (rand() % (right - left + 1));
int pivot = nums[pivot_index];
swap(nums[right], nums[pivot_index]);
int i = left - 1;
int j;
for (j = left; j < right; j++) {
if (nums[j] > pivot) {
swap(nums[j], nums[++i]);
}
}
swap(nums[++i], nums[right]);
return i;
}
};正确的代码
#include <algorithm>
#include <cstdlib>
#include <ctime>
#include <tuple>
#include <vector>
using namespace std;
class Solution {
public:
/**
* @brief 在未排序数组中找到第k大的元素（基于三向切分的快速选择算法）
* @param nums 待搜索的整数数组
* @param k 目标元素的排序位置（第k大）
* @return 第k大的元素值
*
* @note 时间复杂度：平均O(n)，最坏O(n²)（但概率极低）
* @note 空间复杂度：O(1) 原地操作
*/
int findKthLargest(vector<int>& nums, int k) {
int left = 0, right = nums.size() - 1; // 当前搜索范围
while (true) {
// 随机选择基准值避免最坏情况
int pivot_idx = left + rand() % (right - left + 1);
int pivot = nums[pivot_idx];
// 三向切分数组（大于区｜等于区｜小于区）
auto [greater_end, less_start] = threeWayPartition(nums, left, right, pivot);
// 计算各区域元素数量
int greater_cnt = greater_end - left; // 大于区的元素数量
int equal_cnt = less_start - greater_end + 1; // 等于区的元素数量
/* 决策树 */
if (k <= greater_cnt) { // 目标在大于区
right = greater_end - 1; // 缩小搜索范围到左区
} else if (k <= greater_cnt + equal_cnt) { // 命中等于区
return pivot; // 直接返回结果
} else { // 目标在小于区
k -= (greater_cnt + equal_cnt); // 调整k的相对位置
left = less_start + 1; // 缩小搜索范围到右区
}
}
}
private:
/**
* @brief 三向切分数组
* @param nums 待切分数组
* @param left 当前处理区间的左边界
* @param right 当前处理区间的右边界
* @param pivot 基准值
* @return tuple<int, int> 返回两个边界位置：
* - greater_end: 大于区的结束位置（最后一个大于元素的下一位置）
* - less_start: 小于区的开始位置（第一个小于元素的前一位置）
*
* @note 切分结果：
* [left, greater_end) > pivot
* [greater_end, less_start] == pivot
* (less_start, right] < pivot
*/
tuple<int, int> threeWayPartition(vector<int>& nums, int left, int right, int pivot)
{
int greater_end = left; // 大于区的右边界（左侧元素均>pivot）
int less_start = right; // 小于区的左边界（右侧元素均<pivot）
int i = left; // 当前扫描指针
while (i <= less_start) { // 扫描未处理区域
if (nums[i] > pivot) {
// 将大元素交换到大于区末尾
swap(nums[i++], nums[greater_end++]);
} else if (nums[i] == pivot) {
// 等于元素直接跳过
++i;
} else {
// 将小元素交换到小于区头部（注意i不递增）
swap(nums[i], nums[less_start--]);
}
}
return {greater_end, less_start};
}
};3、解题思路
这道题的难度很高，我最开始是采用普通的快速选择，虽然答案是对的，但是会超时，所以我们该用三向切分策略，也就是将数组分为[left, greater_end)、[greater_end, less_start]、(less_start, right] 三个区间，不过需要注意一下开闭区间。
核心思想
-
随机化基准值
每次随机选择基准值 (pivot)，有效避免输入数据有序导致的 O(n²) 最坏时间复杂度。 -
三向切分 (3-Way Partition)
将数组划分为三个区域：- 大于区：所有元素 > pivot
- 等于区：所有元素 == pivot
- 小于区：所有元素 < pivot
-
递归决策
根据各区域元素数量与k值的关系，决定后续搜索方向：- 若k在 大于区：缩小范围到左区间
- 若k在 等于区：直接返回pivot
- 若k在 小于区：调整k值后搜索右区间
执行流程
-
初始化搜索范围
初始处理整个数组 (left=0, right=n-1) -
随机选择基准值
在[left, right]区间随机选取一个元素作为基准值，确保算法鲁棒性 -
三向切分操作
- 使用三个指针：
greater_end：标记大于区的右边界less_start：标记小于区的左边界i：当前扫描指针
- 扫描过程：
- 大元素 → 交换到大于区末尾，指针右移
- 等于元素 → 跳过
- 小元素 → 交换到小于区头部，指针不动（需重新检查）
- 使用三个指针：
- 决策逻辑
if (k <= greater_cnt) { // 目标在左区
right = greater_end - 1;
} else if (k <= greater_cnt + equal_cnt) { // 命中等于区
return pivot;
} else { // 目标在右区
k -= (greater_cnt + equal_cnt);
left = less_start + 1;
}
常见疑问解答
-
为什么用随机pivot？
- 避免最坏时间复杂度，保证算法平均性能
-
less_start指针为何不递增i？
- 交换后的元素来自未处理区域，需重新判断其值
-
如何处理重复元素？
- 三向分区将相同元素集中在中间区域，减少无效比较