LeetCode 703. 数据流中的第 K 大元素

2023-06-05 HGNULB 约 10 分钟

文章信息 9 个分类 6 个标签

分类

题目描述

✅ 703. 数据流中的第 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 大、接纳更有资格的新值，堆的规模仍是 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。

代码实现

import java.util.PriorityQueue;

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();
    }
}

import "container/heap"

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. 最后一块石头的重量	简单	每轮取走两个最大值再放回差值，堆的规模持续收缩而非固定
LCR 059. 数据流中的第 K 大元素	简单	与本题同题，可用来复练构造函数复用 `add` 与空初始数组这两个细节