算法刷题系列（二十四）：堆——动态维护最值与 Top-K

算法刷题系列（二十四）：堆——动态维护最值与 Top-K

如果数据不断到来，而每次都要知道当前最大或最小的元素，反复完整排序通常做了太多工作。堆只维护“堆顶是最值”这一条局部秩序，插入和删除堆顶都只需对数时间。

本文以“保留数据流中最大的 k 个数”为例。为限制空间，维护一个最多 k 项的小顶堆：堆顶是当前保留集合里的最小者，也就是最容易被新元素替换的成员。

一、从 Top-K 目标推导堆的方向

扫描 8, 2, 11, 5, 14, 3，令 k = 3：

1. 先放入 8, 2, 11，小顶堆堆顶为 2；
2. 看到 5，大于堆顶 2，弹出 2、插入 5；
3. 看到 14，替换堆顶 5；
4. 看到 3，小于堆顶 8，不必进入候选集合。

最终保留 {8, 11, 14}。若要求最大的 k 个，使用小顶堆而不是大顶堆；大顶堆会先弹出最大的数，无法快速淘汰保留集合中的最弱者。

二、不变式与正确性

扫描每个元素后，堆中恰好保存“已看过元素里最大的至多 k 个”。堆未满时全部收下；堆满后，只有新值大于最小保留值时，才用它替换堆顶。被替换的旧堆顶不可能属于新的 Top-K，因为已有 k 个数不小于它。

三、Java 实现

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.PriorityQueue;

public class Main {
    static List<Integer> largestK(int[] values, int k) {
        if (k <= 0) {
            return new ArrayList<>();
        }
        PriorityQueue<Integer> heap = new PriorityQueue<>();
        for (int value : values) {
            if (heap.size() < k) {
                heap.offer(value);
            } else if (value > heap.peek()) {
                heap.poll();
                heap.offer(value);
            }
        }
        List<Integer> answer = new ArrayList<>(heap);
        answer.sort(Collections.reverseOrder());
        return answer;
    }

    public static void main(String[] args) {
        assert largestK(new int[]{8, 2, 11, 5, 14, 3}, 3)
                .equals(List.of(14, 11, 8));
        assert largestK(new int[]{4, 4, 1}, 5)
                .equals(List.of(4, 4, 1));
        assert largestK(new int[]{1, 2}, 0).isEmpty();
    }
}

四、Python 实现

import heapq

def largest_k(values: list[int], k: int) -> list[int]:
    if k <= 0:
        return []
    heap: list[int] = []
    for value in values:
        if len(heap) < k:
            heapq.heappush(heap, value)
        elif value > heap[0]:
            heapq.heapreplace(heap, value)
    return sorted(heap, reverse=True)

assert largest_k([8, 2, 11, 5, 14, 3], 3) == [14, 11, 8]
assert largest_k([4, 4, 1], 5) == [4, 4, 1]
assert largest_k([1, 2], 0) == []

五、复杂度、边界与误区

令实际保留元素数为 m = min(k, n)：

• 扫描 n 个数：O(n log m) 时间；
• 堆额外空间：O(m)；
• 若最终按降序输出，再花 O(m log m) 排序。

k <= 0 时返回空结果；k 大于输入长度时返回全部输入并排序；重复值按“元素个数”计入 Top-K。若题意要求前 k 个不同值，应先去重或将堆元素换成不同值的结构。Java PriorityQueue 默认是小顶堆；需要大顶堆时必须显式提供反向比较器。比较器不要用减法实现，可能发生整数溢出。

小结与衔接

堆适合动态维护少量候选的最值；若所有数据已知且只查询一次，排序或选择算法也可能更合适。下一篇将使用小顶堆处理另一类问题：每个有序流只暴露一个当前头部，反复取最小头部即可完成多路归并。