【高频面试题】数组中的第K个最大元素（堆、快排进阶）

数组中的第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<=10^5$
• $-10^4<=nums[i]<=10^4$

解法1（堆维护前k大元素）

时间复杂度$O（nlogk）$ 空间复杂度$O（k)$

class Solution {
public:int findKthLargest(vector<int>& nums, int k) {priority_queue<int, vector<int>, greater<int>> pq; for(auto& num: nums){pq.emplace(num);if(pq.size() > k){pq.pop();   // 堆中元素超过k个，弹出最小的那个}}return pq.top();   }
};

解法2 手写堆维护

思路：

• $nums$种存放二叉堆，索引$[0,n-1]$对应按层序遍历对应的元素，对于下标从0开始的某节点$i$,左右孩子节点编号分别为$i\ast 2+1,i\ast 2+2$
• 下沉操作：$maxHeapify$操作为将二叉堆数组$nums$索引$i$处元素下沉
• 建堆操作：我们从最后一个非叶子节点$heapSize/2-1$开始倒序遍历，从下往上下沉
• 删除操作：每次将堆顶交换到数组末尾，再将$heapSize$减一，最后再调整新的堆顶即可
• 时间复杂度$O（nlogn）$ 空间复杂度$O（logn)$ 因为最大堆需要维护n个结点

class Solution {
public:// down void maxHeapify(vector<int>& a, int i, int heapSize) {int l = i * 2 + 1, r = i * 2 + 2, largest = i;if (l < heapSize && a[l] > a[largest]) {largest = l;}if (r < heapSize && a[r] > a[largest]) {largest = r;}if (largest != i) {swap(a[i], a[largest]);maxHeapify(a, largest, heapSize);}}// initialize void buildMaxHeap(vector<int>& a, int heapSize) {for (int i = heapSize / 2 - 1; i >= 0; i--) {maxHeapify(a, i, heapSize);}}int findKthLargest(vector<int>& nums, int k) {int heapSize = nums.size();buildMaxHeap(nums, heapSize); // 建堆// 执行k-1次提取最大值操作for (int i = nums.size() - 1; i >= nums.size() - k + 1; i--) {swap(nums[0], nums[i]);   // 调整堆顶-- heapSize;              // 删除堆maxHeapify(nums, 0, heapSize); }return nums[0];}
};

解法3（快速选择算法）

我们首先来回顾一下快排的实现

void quickSort(vector<int>& nums, int l, int r) {if (l >= r) return;// - i 初始在左边界左侧（l-1）// - j 初始在右边界右侧（r+1）// - 基准值 x 选中间元素（避免极端情况如全排序数组导致最坏时间复杂度）int i = l - 1, j = r + 1;int x = nums[(l + r) >> 1]; // 位运算代替 (l + r) / 2，等价且更高效// partition ：双指针从两端向中间移动while (i < j) {// 移动左指针 i：跳过所有小于 x 的元素，直到找到 >=x 的元素do i++; while (nums[i] < x); // 移动右指针 j：跳过所有大于 x 的元素，直到找到 <=x 的元素do j--; while (nums[j] > x);//  如果指针未交错，交换左右指针的元素（确保左边 <=x，右边 >=x）if (i < j) swap(nums[i], nums[j]);}// 4. 递归排序左右子数组：quickSort(nums, l, j), quickSort(nums, j + 1, r);
}

• 快速选择（$Quickselect$）算法是一种用于在未排序的列表中找到第k小（或第k大）元素的高效算法。与快速排序一样，它使用分治策略，但不同于快速排序的是，它只递归处理包含目标元素的那一部分，而不是全部。这使得快速选择的平均时间复杂度为$O(n)$，最坏情况为$O(n^2)$，但在实际应用中，通过随机化选取枢轴（$pivot$）可以避免最坏情况。
• 复杂度：递归时，每层时间复杂度为$O(n)$,但并不是都进入左右两部分递归。仅进入一侧递归在平均情况下 数组长度会减半，故平均情况下的时间复杂度为 $n+n/2+n/4+\dots +1=O(n)$

以下是快速选择实现找第K大元素的具体实现：

class Solution {
public:int quick_select(vector<int>& nums, int l, int r, int k) {if (l == r) return nums[k];int i = l - 1, j = r + 1, x = nums[l + r >> 1]; // 若选取nums[l], 极端样例 时间会很久//int x = nums[rand() % (r - l + 1) + l], i = l - 1, j = r + 1; // 随机选取while (i < j) {do i ++ ; while (nums[i] > x); // 注意是第k大 上面的模板是升序排序do j -- ; while (nums[j] < x);if (i < j) swap(nums[i], nums[j]);}if (k <= j) return quick_select(nums, l, j, k);else return quick_select(nums, j + 1, r, k);}int findKthLargest(vector<int>& nums, int k) {//srand(time(0)); // 随机种子return quick_select(nums, 0, nums.size() - 1, k - 1); // 0-base}
};

例题：P1923 【深基9.例4】求第 k 小的数

#include <bits/stdc++.h>using namespace std;const int N = 5e6 + 7;int n, k;
int a[N];int quick_select(int nums[], int l, int r, int k) {if (l == r) return nums[k];int x = nums[l], i = l - 1, j = r + 1;while (i < j) {// paration 方向改一下即可do i ++; while (nums[i] < x); do j --; while (nums[j] > x);if (i < j) swap(nums[i], nums[j]);}if (k <= j) return quick_select(nums, l, j, k);else return quick_select(nums, j + 1, r, k);
}int main() {scanf("%d%d", &n, &k);//getchar();for (int i = 0; i < n; i++) {scanf("%d", &a[i]);}printf("%d\n", quick_select(a, 0, n - 1, k));return 0;
}

参考

• LeetCode官方题解：数组中的第K个最大元素
• yxc常用代码模板1——基础算法
• LeetCode 215. 数组中的第K个最大元素