LeetCode LCR 059. 数据流中的第 K 大元素
题目描述
题意分析
要设计一个类:构造时给定整数
k和一批初始数字,之后每次调用add加入一个新数字,并立即返回加入之后整个数据流中第k大的那个数。这里的「第
k大」是按数值大小排序后的第k个,重复元素各自独立计数,不是去重之后的第k大。比如流里是[5, 5, 4]且k = 2,答案是5而不是4。约束透露的信号很直接:题目保证每次调用
add之后流中至少有k个元素,所以不存在「元素不够k个该返回什么」的未定义情况;同时add的调用次数可以远大于k,这意味着每次都把全部历史元素排序是纯粹的浪费。更关键的观察藏在问题本身:这个数据流只增不减。一个当前排在第
k名之后的数字,想要成为第k大,必须先有元素从流中消失,而这永远不会发生。所以历史里绝大多数数字都可以立刻丢掉,真正需要长期保留的只有「当前最大的那k个」,答案就是这一小撮里最小的那个。边界上要留意三点:
k = 1时答案退化成全局最大值;初始数组允许为空,构造函数不能假设它至少有一个元素;元素值可以是负数,不能拿0或-1当作「还没有元素」的哨兵。
解法:固定容量小根堆
核心思路
暴力做法是把所有见过的数字存进一个数组,每次
add之后整体排序,取倒数第k个。逻辑上完全正确,但单次add就要 $O(n \log n)$,n次调用累计到 $O(n^2 \log n)$,在调用量上万时必然超时。瓶颈在于每次都对整个历史重新排序,而其中绝大部分元素根本没有资格成为答案。既然流只增不减,一个已经掉出前
k名的数字就永久失去了翻身机会,为它排序是白费力气。于是可以只保留最大的
k个数。这批数里最小的那个恰好就是全局第k大——比它大的正好有k - 1个,比它小的全在被丢弃的那堆里。所以需要的容器必须支持两件事:快速拿到「这批数里的最小值」,以及快速把最小值换掉。小根堆正好同时满足,堆顶就是最小值,弹出与插入都是 $O(\log k)$。维护的不变量是:每次
add返回之前,堆中恰好保存着当前数据流中最大的 $\min(k,\ \text{流中元素个数})$ 个元素,且堆顶是它们之中的最小值。因为题目保证返回时流中至少有k个元素,这个不变量落到返回时刻就等价于「堆中正好是最大的k个,堆顶即第k大」。新元素到来时如何维持不变量,只有两种情况:堆还没装满
k个,那它无条件属于「最大的若干个」,直接放进去;堆已满时,它只有严格大于堆顶才有资格挤进前k大,此时弹掉堆顶再放入,规模仍是k;否则它连当前第k大都比不过,直接丢弃即可。
解题步骤
- 构造函数里把
k存成成员变量,并建一个空的小根堆。堆必须是小根而不是大根:需要被随时替换掉的是这批数里最弱的那个,只有小根堆能 $O(1)$ 看到它。- 把初始数组里的数字逐个交给
add处理,而不是另写一套插入逻辑。复用的好处是「未满就放、满了就比堆顶」这条规则只在一处实现,初始数组长度小于k、等于k、大于k三种情况自动被同一份代码覆盖。add的第一分支:堆的大小还不足k,直接入堆。此时不做任何比较,因为元素总数还没超过k,每个来的数都属于「最大的若干个」。add的第二分支:堆已满且新值严格大于堆顶。先弹出堆顶再压入新值,顺序不能颠倒——先压后弹会让堆一度有k + 1个元素,虽然本题中弹出的仍是最小值、结果碰巧相同,但堆规模超标的写法在换成「保留最小的 k 个」等变体时会立刻出错,不如始终保持规模恒定。add的第三分支(隐含):堆已满且新值不大于堆顶,什么都不做。丢弃是安全的,因为它连当前第k大都比不上,而未来只会有更多元素加入,它的排名只会更靠后。- 返回堆顶。注意返回的是查看而非弹出,答案要留在堆里继续参与后续比较。
以
k = 3、nums = [4, 5, 8, 2],随后依次add(3)、add(5)、add(10)、add(9)、add(4)走一遍:构造阶段,4入堆,堆为{4};5入堆,堆为{4, 5};8入堆,堆为{4, 5, 8},此时已满,堆顶为4;2到来,堆已满且2 < 4,直接丢弃,堆仍是{4, 5, 8}。构造结束。add(3):堆满且3 < 4,丢弃,返回堆顶4。add(5):5 > 4,弹出4压入5,堆变成{5, 5, 8},返回堆顶5——注意这里两个5同时存在,重复元素各自计数,符合题意。add(10):10 > 5,弹出一个5压入10,堆为{5, 8, 10},返回5。add(9):9 > 5,弹出5压入9,堆为{8, 9, 10},返回8。add(4):堆满且4 < 8,丢弃,返回8。整个过程中被丢掉的2、3、4从未影响过任何一次答案,而堆的规模始终没有超过3。
代码实现
class KthLargest {
private final int k;
private final PriorityQueue<Integer> heap;
public KthLargest(int k, int[] nums) {
this.k = k;
heap = new PriorityQueue<>();
for (int num : nums) {
add(num);
}
}
public int add(int val) {
if (heap.size() < k) {
heap.offer(val);
} else if (val > heap.peek()) {
// 堆中只保留当前最大的 k 个数,堆顶就是第 k 大。
heap.poll();
heap.offer(val);
}
return heap.peek();
}
}
type IntHeap struct {
data []int
}
func (h IntHeap) Len() int {
return len(h.data)
}
func (h IntHeap) Less(i int, j int) bool {
return h.data[i] < h.data[j]
}
func (h IntHeap) Swap(i int, j int) {
h.data[i], h.data[j] = h.data[j], h.data[i]
}
func (h *IntHeap) Push(x any) {
h.data = append(h.data, x.(int))
}
func (h *IntHeap) Pop() any {
val := h.data[len(h.data)-1]
h.data = h.data[:len(h.data)-1]
return val
}
type KthLargest struct {
k int
heap IntHeap
}
func Constructor(k int, nums []int) KthLargest {
kth := KthLargest{k: k, heap: IntHeap{}}
heap.Init(&kth.heap)
for _, num := range nums {
kth.Add(num)
}
return kth
}
func (this *KthLargest) Add(val int) int {
if this.heap.Len() < this.k {
heap.Push(&this.heap, val)
} else if val > this.heap.data[0] {
// 堆中只保留当前最大的 k 个数,堆顶就是第 k 大。
heap.Pop(&this.heap)
heap.Push(&this.heap, val)
}
return this.heap.data[0]
}
复杂度分析
- 时间复杂度:构造函数 $O(n \log k)$,单次
add为 $O(\log k)$。凭什么:堆的规模被显式限制在k以内,一次插入或弹出只沿着树高走 $O(\log k)$ 步;构造函数把n个初始元素各过一遍这套逻辑,而每次add最多触发一次弹出加一次插入,与已经流过多少元素无关。- 空间复杂度:$O(k)$。凭什么:容器里任何时刻至多有
k个元素,被丢弃的数字不占用任何存储,所以内存与数据流的总长度无关,只与k有关。
关键点总结
- 「第
k大」可以改写成「最大的k个数里的最小值」。这个改写是全题的支点:它把一个需要全局排序的查询,变成了对一个固定容量集合的最小值查询。凡是遇到「第 K 大 / 第 K 小 / 前 K 个」,先做这一步等价改写往往就能看到解法。- 求第
k大用小根堆,求第k小用大根堆——方向和直觉相反。原因是堆顶必须是「最容易被淘汰的那个」,求最大的k个时最容易被淘汰的正是其中最小者。- 数据流只增不减这一性质是丢弃元素的许可证。如果题目改成支持删除,被丢掉的数字就可能重新变成答案,这套做法立刻失效,得换成有序集合或者对顶堆。
- 容量固定的堆把复杂度里的
n换成了k。当k远小于流长度时收益巨大,这也是「维护规模上界」这一类优化的通用价值。- 面试视角:先说暴力排序的 $O(n \log n)$ 单次代价,再点出「只增不减 ⇒ 掉出前
k名的数永远回不来」,最后才引出小根堆,这条推导链比直接甩出答案更能拿分。面试官真正想听的是你为什么敢丢数据。- 面试视角:常见追问有两个。一是「如果要求第
k小怎么改」,答堆的方向反过来、比较条件改成小于堆顶;二是「如果还要支持删除任意元素」,答需要用可删除的有序结构,或者用两个堆加延迟删除来维护。能主动区分「静态数组求第 K 大」(快速选择 $O(n)$ 更优)和「数据流求第 K 大」(必须用堆)是加分项。
易错点总结
- 错误写法:用大根堆保留最大的
k个数。用例k = 3,流为[4, 5, 8, 2]→ 大根堆的堆顶是8,返回8,正确答案是4;堆顶必须是这批数里最小的那个才等于第k大。- 错误写法:把所有元素都塞进小根堆不做容量限制。用例
k = 3,流为[4, 5, 8, 2]→ 堆顶是全局最小值2,返回2,正确答案是4;同时内存随流长度无限增长。- 错误写法:堆满时不比较大小,一律弹出堆顶再压入新值。用例
k = 3,堆为{4, 5, 8}时add(2)→ 弹出4压入2,堆变成{2, 5, 8},返回2,正确答案仍是4;一次错误替换会永久污染后续所有查询。- 错误写法:未满判断写成
heap.size() <= k。用例k = 3,流为[4, 5, 8, 2]→ 堆被装到4个元素{2, 4, 5, 8},堆顶变成2,返回2,正确答案是4;容量上界写错一格,堆顶含义就整体错位一名。- 错误写法:在构造函数里另写一套「先全部入堆再截断」的逻辑。用例
k = 3,nums = []→ 空数组上取堆顶或做截断时访问不存在的元素,抛空指针或下标越界异常。- 错误写法:
add里返回前用poll取堆顶而不是peek。用例k = 3,堆为{4, 5, 8},连续两次add(2)→ 第一次返回4但把它弹走了,堆只剩两个元素,第二次2被当作未满直接入堆,返回2,正确答案是4。- 错误写法:用
0或-1初始化答案,认为堆空时返回它。用例k = 1,流为[-5]→ 返回0或-1,正确答案是-5;元素允许为负,任何常数哨兵都可能与真实值冲突。- 错误写法:把「第
k大」理解成去重后的第k大,入堆前先判重。用例k = 2,流为[5, 5, 4]→ 第二个5被当作重复丢弃,堆里只有{4, 5},返回4,正确答案是5。
相似题目
| 题目 | 难度 | 考察点 |
|---|---|---|
| 215. 数组中的第K个最大元素 | 中等 | 数组静态给定且只查一次,快速选择的 $O(n)$ 期望优于堆,可对比取舍 |
| 239. 滑动窗口最大值 | 困难 | 元素会因窗口移出而失效,不能只增不减,需要单调队列或堆配合延迟删除 |
| 295. 数据流的中位数 | 困难 | 查询位置随流长度变化,要用大小两个堆对顶并动态调整平衡 |
| 347. 前 K 个高频元素 | 中等 | 比较键是出现次数而非元素本身,需先统计频次再对频次做定容筛选 |
| 355. 设计推特 | 中等 | 从多条有序推文流中合并取最新 10 条,是多路归并版本的 Top K |
| 378. 有序矩阵中第 K 小的元素 | 中等 | 数据自带行列有序性,可用二分答案计数,复杂度低于无脑堆 |
| 973. 最接近原点的 K 个点 | 中等 | 求最小的 k 个,堆的方向要反过来用大根堆,是本题的镜像 |
| 1046. 最后一块石头的重量 | 简单 | 每轮取走两个最大值再放回差值,堆的规模持续收缩而非固定 |