LeetCode 215. 数组中的第 K 个最大元素
本节目标
用容量为 k 的小根堆维护最大的 k 个元素,理解堆顶作为淘汰线的含义。
这道题是堆与优先队列解题框架的第一道母题。目标不是把数组完全排好序,而是在遍历过程中始终保留最有希望成为答案的 k 个元素。
题意与约束
给定整数数组 nums 和整数 k,返回数组排序后第 k 个最大的元素。这里求的是有序位置,而不是第 k 个不同的数,因此重复元素需要正常计数。
例如 nums = [3, 2, 1, 5, 6, 4]、k = 2,降序排列为 [6, 5, 4, 3, 2, 1],答案是 5。
第一反应:排序能做,但做多了
将整个数组排序后直接读取对应位置,时间复杂度为 O(n log n)。这个方法正确,也常适合作为基准解,但它维护了所有元素之间的完整顺序。
本题只需要一个边界:最大的 k 个元素中最小的那个。我们可以只维护这 k 个候选,而不关心其他元素如何排序。
固定容量小根堆
建立一个容量不超过 k 的小根堆 top:
- 遍历一个数就先加入堆;
- 若堆大小超过
k,删除堆顶的最小值; - 遍历结束时,堆中保留全局最大的
k个数; - 堆顶是这
k个数中最小的数,因此正是全局第k大。
处理 nums = [3, 2, 1, 5, 6, 4]、k = 2 时:
| 读入元素 | 调整后的堆内元素 | 堆顶含义 |
|---|---|---|
| 3 | [3] | 当前候选最小值 |
| 2 | [2, 3] | 当前第 2 大 |
| 1 | [2, 3] | 1 被淘汰 |
| 5 | [3, 5] | 2 被淘汰 |
| 6 | [5, 6] | 3 被淘汰 |
| 4 | [5, 6] | 4 被淘汰 |
每轮结束后都有一个重要不变量:堆中保存“已处理元素里最大的至多 k 个数”。这也是算法正确性的核心。
为什么不是大根堆
大根堆可以装入全部元素,再连续弹出 k - 1 次,最后读取堆顶。它的空间复杂度是 O(n)。
小根堆方案让无资格进入前 k 名的元素尽早离开,堆的规模始终不超过 k,空间复杂度降为 O(k)。求“最大的若干个”时,用小根堆维护淘汰线;求“最小的若干个”时,则反过来用大根堆。
代码实现
C++ 的 priority_queue 默认是大根堆,因此需要通过 greater<int> 声明小根堆;Python 的 heapq 默认就是小根堆。两份实现都直接维护同一个固定容量不变量。
- C++
- Python
#include <functional>
#include <queue>
#include <vector>
using namespace std;
class Solution {
public:
int findKthLargest(vector<int>& nums, int k) {
priority_queue<int, vector<int>, greater<int>> top;
for (int value : nums) {
top.push(value);
if (static_cast<int>(top.size()) > k) {
top.pop();
}
}
return top.top();
}
};
import heapq
class Solution:
def findKthLargest(self, nums, k):
top = []
for value in nums:
heapq.heappush(top, value)
if len(top) > k:
heapq.heappop(top)
return top[0]
复杂度分析
- 时间复杂度:
O(n log k)。每个元素最多进行一次插入和一次删除,堆大小不超过k。 - 空间复杂度:
O(k)。堆中最多保存k个元素。
当 k 远小于 n 时,这个方法比完整排序更能体现“只维护必要信息”的思想。
易错点
- 把“第
k大”误解为“第k个不同的数”,从而错误去重。 - 使用容量为
k的大根堆,堆顶会是候选中的最大值,不能直接得到第k大。 - 在堆大小达到
k时就弹出。应当在大小严格超过k时才淘汰。 - 最后把堆中任意元素当作答案。答案必须是小根堆堆顶。
模式迁移
固定容量堆的本质是“只保留当前最优的 k 个候选”。后续可以将候选从数值扩展为:
- 元素与出现次数,解决前 K 高频问题;
- 距离与点,解决离原点最近的 K 个点;
- 任务收益与截止时间,解决部分贪心调度问题。
每次都要先确认比较依据:究竟按元素值、频率、距离,还是其他优先级决定淘汰谁。