如何用堆维护数据流中的第 K 大元素？

简化版

用一个大小为 K 的小顶堆维护当前最大的 K 个元素。堆顶就是这 K 个元素里最小的，也就是整个数据流目前的第 K 大；新元素比堆顶大才有资格进入堆。

详细版

数据流第 K 大和一次性数组 Top K 很像，但它的重点是“边到边查询”：每次 add(x) 后都要马上返回当前第 K 大。

核心做法：

维护一个小顶堆 heap。
堆大小小于 K 时，直接放入新元素。
堆大小等于 K 时，只在 x > heap.peek() 时替换堆顶。
堆顶 heap.peek() 始终是当前第 K 大。

复杂度：每次插入最多一次 poll + offer，时间 O(log K)；堆里最多 K 个元素，空间 O(K)。如果数据流还不足 K 个，通常按题意返回当前堆顶、空值或等待到满 K 后再返回。

class KthLargest {
    private final int k;
    private final PriorityQueue<Integer> heap = new PriorityQueue<>();

    KthLargest(int k, int[] nums) {
        this.k = k;
        for (int x : nums) add(x);
    }

    int add(int val) {
        if (heap.size() < k) heap.offer(val);
        else if (val > heap.peek()) {
            heap.poll();
            heap.offer(val);
        }
        return heap.peek();
    }
}

完整版教学

一、这题为什么不是每次排序

数据流的特点是元素不断到来，查询也不断发生。假设已经来了 n=1,000,000 个元素，每次新增后都全量排序，单次代价接近 O(n log n)，连续查询会非常浪费。第 K 大并不要求你知道所有元素的完整顺序，只需要知道“当前最大的 K 个元素”以及其中最小的那一个。

堆正好适合维护这个边界。小顶堆的堆顶是堆内最小值，如果堆里放的是当前最大的 K 个元素，那么堆顶自然就是第 K 大。这个设计把“大问题的全排序”变成“维护 K 个候选者”。

目标：第 3 大
当前最大的 3 个： [10, 9, 8]
小顶堆堆顶：8  -> 当前第 3 大
其他元素：7, 5, 2  都不影响答案

二、为什么求第 K 大要用小顶堆

这道题最容易反直觉：求“大”的东西，却用“小顶堆”。原因是堆顶承担的是“淘汰门槛”，不是最终最大值。维护最大的 K 个元素时，最该被淘汰的是这 K 个候选者中最小的那个，所以要让它能被 O(1) 看到。

例如 K=3，堆里是 [8, 10, 9]，堆顶为 8。新元素 7 连门槛都达不到，丢弃；新元素 11 超过门槛，说明它应该进入 Top 3，于是淘汰 8。如果用大顶堆，堆顶是 10 或 11，你反而拿不到“当前最弱候选者”，替换判断会变复杂。

需求	保留集合	合适堆	堆顶含义
第 K 大 / 最大 K 个	当前最大的 K 个	小顶堆	第 K 大，淘汰门槛
第 K 小 / 最小 K 个	当前最小的 K 个	大顶堆	第 K 小，淘汰门槛
只要最大值	全部候选	大顶堆	最大值

三、用一个数字例子走完整过程

设 K=3，初始数据流依次到来：4, 5, 8, 2, 10, 9。前 3 个元素先填满堆，之后每个元素都和堆顶比较。堆内部数组顺序不代表排序结果，只要堆顶正确即可。

add 4: heap=[4]          不足 3 个
add 5: heap=[4,5]        不足 3 个
add 8: heap=[4,5,8]      第 3 大 = 4
add 2: 2 <= 4, 丢弃      第 3 大 = 4
add 10: 10 > 4, 替换     heap 存 {5,8,10}, 第 3 大 = 5
add 9: 9 > 5, 替换       heap 存 {8,9,10}, 第 3 大 = 8

这个过程说明堆不是保存所有历史数据，而是保存“仍可能影响第 K 大答案”的候选集合。小于等于堆顶的元素不可能进入当前最大的 K 个，直接丢弃是安全的。

四、代码不变式与边界处理

实现时要始终维护两个不变式：第一，堆大小不超过 K；第二，堆里保存的是当前最大的 K 个元素。只要这两个条件成立，堆顶就是当前第 K 大。判断逻辑可以写成“先入堆再裁剪”，也可以写成“满了才比较”，两种都可以。

int add(int x) {
    heap.offer(x);
    if (heap.size() > k) {
        heap.poll(); // 弹掉候选集合中最小的
    }
    return heap.peek();
}

这种写法更短，但会把明显不够格的元素也先插入一次，仍然是 O(log K)。如果 K=0 通常是非法输入，应在构造时拒绝；如果数据量暂时小于 K，要按题目约定处理，不能默认存在第 K 大。

五、复杂度为什么是 O(log K) 而不是 O(log n)

堆操作的复杂度取决于堆大小。这里堆最多只有 K 个元素，所以每次 offer 或 poll 是 O(log K)，不是 O(log n)。当 K=100、数据流来了 n=10^9 个元素时，单次操作仍然只和 100 的对数相关。

每个元素最多：
  比较堆顶：O(1)
  插入或替换：O(log K)
总处理 n 个元素：O(n log K)
额外空间：O(K)

和快速选择相比，堆法适合流式数据，因为它不需要拿到完整数组，也不会修改原数据。快速选择平均 O(n) 很强，但它是一次性离线算法，不适合每来一个元素就返回答案。

六、常见误区与追问

记忆钩子：第 K 大维护“最大的 K 个”，堆顶不是最大，而是进入 Top K 的门槛。

误区：求第 K 大应该用大顶堆。 大顶堆能快速拿最大值，但不能快速找到 Top K 里最该淘汰的最小候选者。
误区：堆里的数组顺序就是从小到大排序。 堆只保证父子关系，不能把内部数组当有序数组使用。
误区：每次 add 都必须保存所有历史元素。 只要题目只问第 K 大，保存 K 个候选者就够了。
追问：新元素等于堆顶要不要替换？ 不需要，替换后第 K 大不变；除非题目要求保留元素身份或稳定性。
追问：数据不足 K 个怎么办？ 这取决于接口定义，可以返回当前最小值、空值或等到满 K 后再提供答案。
追问：为什么复杂度是 O(log K)？ 因为堆大小被限制为 K，堆高约为 log K。

七、加强记忆

数据流第 K 大的核心不是“排序”，而是“守住门槛”。把当前最大的 K 个元素放进小顶堆，堆顶就是这 K 个候选者中最弱的一个，也就是第 K 大。新元素只有超过堆顶才会改变候选集合，否则直接丢弃。记忆时抓住三句话：保留 K 个候选者，小顶堆顶当门槛，每次更新 O(log K)。这样无论题目问第 K 大、实时排行榜还是流式分数线，都能自然想到同一套模型。

简化版

详细版