当前位置：首页 > article >正文

LeetCode-215. 数组中的第K个最大元素

article 2025/3/25 21:17:21

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