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

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

链接 #

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

题目 #

给定整数数组 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

提示：

  • 1 <= k <= nums.length <= 105
  • -104 <= nums[i] <= 104

解答 #

#include <vector>
using namespace std;

class Solution {
public:
    void quickSelect(vector<int>& nums, int l, int r, int k) {
        if (l >= r) return;
        
        int pivot = nums[l + (r - l) / 2];
        int i = l, j = r, e = l;
        
        // 三路划分过程
        while (e <= j) {
            if (nums[e] < pivot) {
                swap(nums[e++], nums[i++]);
            } else if (nums[e] > pivot) {
                swap(nums[e], nums[j--]);
            } else {
                e++;
            }
        }
        
        // 递归调用
        if (k < i) quickSelect(nums, l, i - 1, k);
        else if (k > j) quickSelect(nums, j + 1, r, k);
    }

    int findKthLargest(vector<int>& nums, int k) {
        int n = nums.size();
        quickSelect(nums, 0, n - 1, n - k);
        return nums[n - k];
    }
};

每次经过划分后，我们都会得到三个区间：[l, i-1]、[i, j]和[j+1, r]。并且能确定 [i, j] 之间的元素的位置已经是最终位置。

quickSelect 函数的作用是让位于索引 k 的元素为其排序后的最终正确位置。所以我们实际上取 n - k 作为索引，由于排序后的数组递增，得到的是第 k 大的元素。

再看递归调用的两行，当 k < i，就意味着索引位于 k 的最终元素需要在左侧区间确定。k > j 同理。

另外再解释一下三路划分的循环过程：

  • i, j 是左右双指针
  • e 是当前遍历的指针，会从左到右一步步移动。
    • 如果 nums[e] < pivot，就将 nums[e] 与 nums[i] 交换，然后 i 和 e 都向右移动一步。
    • 也即：如果当前数小于基准数，就将当前数换到左边区间，同时左边区间的右边界向右移动一位。而当前指针也要移动，这是因为左边区间换来的数都被 e 遍历过，一定是小于基准数的。
    • 如果 nums[e] > pivot，就将 nums[e] 与 nums[j] 交换，然后 j 向左移动一步。
    • 也即：如果当前数大于基准数，就将当前数换到右边区间，同时右边区间的左边界向左移动一位。当前指针不移动，因为换来的数没有被 e 处理过，依然要再检查。（比如换来的可能是一个大于基准数的数）
    • 如果 nums[e] == pivot，就将 e 向右移动一步。这种情况下，当前数已经在正确的位置上了，不需要换。

可根据下面的幻灯片理解：

Edit Edit this page