[LC]215.数组中的第k个最大元素

题目

给定整数数组 nums 和整数 k，请返回数组中第 k 个最大的元素。

请注意，你需要找的是数组排序后的第 k 个最大的元素，而不是第 k 个不同的元素。

你必须设计并实现时间复杂度为 O(n) 的算法解决此问题。

分析

快速选择

经典算法，寻找第k个元素。可以借用快速排序的划分操作，每次划分操作将一个枢轴元素放置到它排好序后的正确位置上，且所有小于等于枢轴元素的元素都在该位置之前，所有大于等于枢轴元素的元素都在该位置之后。如果这个位置等于k，则我们直接返回该位置上的值即可；如果这个位置大于k，说明我们要找到的元素一定在枢轴元素之前，递归的求解一个更小的子问题即可；如果这个位置小于k，同理。

代码

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        return findKthElement(nums, 0, nums.size() - 1, nums.size() - k);
    }
    int findKthElement(vector<int>& nums, int low, int high, int k)
    {
        int pivot_idx = partition(nums, low, high);
        if (pivot_idx == k) return nums[k];
        else if (pivot_idx > k) return findKthElement(nums, low, pivot_idx - 1, k);
        else return findKthElement(nums, pivot_idx + 1, high, k);
    }
    int partition(vector<int>& nums, int low, int high)
    {
        int pivot = nums[low];
        int l = low + 1, r = high;
        while (true)
        {
            while (l <= r && nums[l] < pivot) ++l;
            while (l <= r && nums[r] > pivot) --r;
            if (l >= r) break;
            swap(nums[l], nums[r]);
            ++l; --r;
        }
        swap(nums[low], nums[r]);
        return r;
    }
};

注意代码中的partition算法，该算法将区间$[low,high]$划分为两个区间，一个区间是$[low + 1, l)$，此区间满足任意元素都严格小于pivot，另一个区间是$(r, high]$，此区间满足任意元素都严格大于pivot。当跳出循环的时候，有两种可能：

1. $l = r$，此时$nums[l] >= pivot$且$nums[l] <= pivot$。即$nums[l] = nums[r] = pivot$.
2. $l = r + 1$, 此时$r < l$，因此$nums[r] < pivot, nums[l] > pivot$.因此$r$是枢轴位置。

注意条件中的$l <= r$的等号不能省略，这是为了确保当$l == r$的时候，仍然会对该位置的元素进行检测，使之满足区间约束条件。反之，如果我们省略等号，那么当$l == r$的时候，有可能因为$l <r$这个条件不成立导致循环提前终止，而没有对当前位置的元素进行检测，其有可能不符合区间的约束条件。

大顶堆

我们将元素维护在一个大顶堆中，然后依次取出的第k个元素即为我们所求。

代码

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        int n = nums.size();
        buildHeap(nums, n);
        for (int i = 1; i <= k; ++i)
        {
            swap(nums[0], nums[n - i]);
            heapAdjust(nums, 0, n - i);
        }
        return nums[n - k];
    }
    void buildHeap(vector<int> &nums, int n)
    {
        for (int i = n / 2 - 1; i >= 0; --i)
        {
            heapAdjust(nums, i, n);
        }
    }
    void heapAdjust(vector<int> &nums, int k, int n)
    {
        int i = left(k);
        while (i < n)
        {
            if (i + 1 < n && nums[i + 1] > nums[i])++i;
            if (nums[k] >= nums[i]) break;
            swap(nums[k], nums[i]);
            k = i;
            i = left(k);
        }
    }
    int left(int i){ return 2 * i + 1; }
};

// c++标准库算法实现
class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        int n = nums.size();
        make_heap(nums.begin(), nums.end());
        for (int i = 0; i < k; ++i)
        {
            pop_heap(nums.begin(), nums.end() - i);
        }
        return nums[n - k];
    }
};

时间复杂度

建堆的时间复杂度是$O(n)$，每一次删除一个元素的时间复杂度是$O(logn)$，因此总的时间复杂度是$O(n + klogn).$

空间复杂度

$O(1).$