第23篇-堆与优先队列-TopK问题的常规武器

概述

上一篇我们学习了二叉搜索树，核心是利用"有序"来加速查找、插入、删除和验证。

这一篇我们继续学习一种同样很常见的数据结构：堆。

堆在算法题里最常见的用途是：

维护前 K 大
维护前 K 小
动态取最值
合并多个有序序列
处理实时流式数据

很多题的题面都会写成：

text 复制代码

给你一组数据，返回最大的 K 个数

或者：

text 复制代码

返回出现频率最高的 K 个元素

这类题如果每次都排序，往往会多做很多无用功。

堆的价值就在于：

text 复制代码

只保留最有用的那一小部分元素

学完这篇，你应该能判断什么时候该用堆，并能写出 TopK、动态最值和优先队列的标准模板。

核心概念：堆到底是什么

堆是一种特殊的完全二叉树结构。

它最重要的性质是：

小根堆：父节点值小于等于子节点值
大根堆：父节点值大于等于子节点值

小根堆示意

text 复制代码

在小根堆中，根节点是当前最小值。

大根堆示意

text 复制代码

在大根堆中，根节点是当前最大值。

堆和排序的区别

堆并不保证整棵树完全有序。

它只保证：

text 复制代码

根节点一定是当前最值

所以堆非常适合"只关心最值"的问题，而不适合一次性拿到全部排序结果。

堆不是全局有序结构，它只保证根节点永远是当前最小值或最大值。

优先队列：Java 里怎么用堆

Java 中，最常用的堆实现是 PriorityQueue。

默认情况下，PriorityQueue 是小根堆。

小根堆写法

java 复制代码

import java.util.PriorityQueue;

PriorityQueue<Integer> pq = new PriorityQueue<>();
pq.offer(3);
pq.offer(1);
pq.offer(2);

System.out.println(pq.poll()); // 1
System.out.println(pq.poll()); // 2
System.out.println(pq.poll()); // 3

大根堆写法

java 复制代码

import java.util.Collections;
import java.util.PriorityQueue;

PriorityQueue<Integer> pq = new PriorityQueue<>(Collections.reverseOrder());
pq.offer(3);
pq.offer(1);
pq.offer(2);

System.out.println(pq.poll()); // 3
System.out.println(pq.poll()); // 2
System.out.println(pq.poll()); // 1

常用 API

API	作用
`offer(x)`	插入元素
`poll()`	弹出堆顶元素
`peek()`	查看堆顶元素
`size()`	当前元素个数

为什么 `PriorityQueue` 适合刷题

因为它提供了：

插入
删除堆顶
查询堆顶

这些操作的复杂度一般都是：

text 复制代码

O(log n)

Java 的 PriorityQueue 默认是小根堆，能高效维护当前最值。

题型一：数组中的前 K 大元素

题目：

给定一个整数数组 nums 和整数 k，返回数组中最大的 k 个元素。

这是堆的经典入门题。

解题思路

如果直接排序，复杂度是：

text 复制代码

O(n log n)

但其实我们只关心最大的 k 个元素，不需要把整个数组完全排好序。

更好的做法是维护一个大小为 k 的小根堆：

先把前 k 个数放进堆
之后每来一个新数，就和堆顶比较
如果新数更大，就替换堆顶
最后堆里剩下的就是最大的 k 个数

为什么用小根堆？

因为堆顶保存的是当前这 k 个数里最小的那个。

一旦来了更大的数，说明它有资格进入 TopK。

Java 代码实现

java 复制代码

import java.util.Arrays;
import java.util.PriorityQueue;

class Solution {
    public int[] topKFrequent(int[] nums, int k) {
        PriorityQueue<Integer> pq = new PriorityQueue<>();

        for (int num : nums) {
            if (pq.size() < k) {
                pq.offer(num);
            } else if (num > pq.peek()) {
                pq.poll();
                pq.offer(num);
            }
        }

        int[] ans = new int[k];
        for (int i = k - 1; i >= 0; i--) {
            ans[i] = pq.poll();
        }

        return ans;
    }
}

为什么最后要倒序取出

小根堆弹出来的是从小到大。

如果题目要求返回大的在前面，就可以倒序填充数组。

维护一个大小为 k 的小根堆，堆顶永远是当前 TopK 里最小的元素。

题型二：数组中的前 K 小元素

前 K 小和前 K 大是完全对称的。

只不过这次通常维护一个大小为 k 的大根堆。

解题思路

先把前 k 个数放进大根堆
后续元素如果比堆顶更小，就替换堆顶
最后堆里保存的就是最小的 k 个数

Java 代码实现

java 复制代码

import java.util.Collections;
import java.util.PriorityQueue;

class Solution {
    public int[] smallestK(int[] nums, int k) {
        PriorityQueue<Integer> pq = new PriorityQueue<>(Collections.reverseOrder());

        for (int num : nums) {
            if (pq.size() < k) {
                pq.offer(num);
            } else if (num < pq.peek()) {
                pq.poll();
                pq.offer(num);
            }
        }

        int[] ans = new int[k];
        for (int i = 0; i < k; i++) {
            ans[i] = pq.poll();
        }

        return ans;
    }
}

为什么前 K 小用大根堆

因为我们要保留的是最小的 k 个元素。

堆顶应该是这 k 个元素里最大的那个，方便判断新元素能不能挤进去。

维护一个大小为 k 的大根堆，堆顶永远是当前 TopK 里最大的元素。

题型三：出现频率最高的 K 个元素

这是堆最常见的 TopK 变形题之一。

题目：

给定一个数组，返回出现频率最高的 k 个元素。

解题思路

这类题的关键是两步：

先统计频率
再按频率维护 TopK

我们可以用哈希表统计每个数出现的次数，然后用小根堆保留频率最高的 k 个元素。

Java 代码实现

java 复制代码

import java.util.HashMap;
import java.util.Map;
import java.util.PriorityQueue;

class Solution {
    public int[] topKFrequent(int[] nums, int k) {
        Map<Integer, Integer> freq = new HashMap<>();
        for (int num : nums) {
            freq.put(num, freq.getOrDefault(num, 0) + 1);
        }

        PriorityQueue<int[]> pq = new PriorityQueue<>((a, b) -> a[1] - b[1]);

        for (Map.Entry<Integer, Integer> entry : freq.entrySet()) {
            int num = entry.getKey();
            int count = entry.getValue();

            if (pq.size() < k) {
                pq.offer(new int[]{num, count});
            } else if (count > pq.peek()[1]) {
                pq.poll();
                pq.offer(new int[]{num, count});
            }
        }

        int[] ans = new int[k];
        for (int i = 0; i < k; i++) {
            ans[i] = pq.poll()[0];
        }

        return ans;
    }
}

为什么堆里存的是数组

因为我们需要同时保存：

元素值
频率

所以堆节点通常写成：

java 复制代码

new int[]{num, count}

比较器只按 count 排序即可。

复杂度分析

统计频率：O(n)
堆操作：O(m log k)，m 是不同元素个数

总体复杂度通常写成：

text 复制代码

O(n log k)

先用哈希表统计，再用小根堆保留频率最高的 k 个元素。

题型四：合并 K 个有序数组或链表

堆除了做 TopK，也常用于合并多个有序序列。

题目通常是：

给定 k 个升序链表，合并成一个升序链表。

或者：

给定 k 个升序数组，合并成一个升序结果。

这类题的核心都是：

text 复制代码

每次取当前所有序列中的最小头部元素

所以很自然地想到小根堆。

合并 K 个升序链表示例

java 复制代码

import java.util.PriorityQueue;

class Solution {
    public ListNode mergeKLists(ListNode[] lists) {
        PriorityQueue<ListNode> pq = new PriorityQueue<>((a, b) -> a.val - b.val);

        for (ListNode node : lists) {
            if (node != null) {
                pq.offer(node);
            }
        }

        ListNode dummy = new ListNode(0);
        ListNode tail = dummy;

        while (!pq.isEmpty()) {
            ListNode node = pq.poll();
            tail.next = node;
            tail = node;

            if (node.next != null) {
                pq.offer(node.next);
            }
        }

        return dummy.next;
    }
}

为什么每个链表只放头节点

因为每个链表都是有序的。

当前头节点弹出后，下一个可能成为最小值的，只会是它后面的节点。

这就像归并排序里的多路归并，只是这里用堆来维护最小头部。

每次从多个有序序列里取当前最小头部，小根堆是最自然的工具。

题型五：动态维护中位数

这是堆的一个非常经典扩展题。

题目：

设计一个数据结构，支持不断插入数字，并随时查询当前中位数。

这类题通常需要两个堆：

大根堆：保存较小的一半
小根堆：保存较大的一半

核心思想

保持两个堆的数量差不超过 1：

大根堆的堆顶是左半部分最大值
小根堆的堆顶是右半部分最小值

这样：

总数为奇数时，中位数就是较大那个堆的堆顶
总数为偶数时，中位数就是两个堆顶的平均值

代码示意

java 复制代码

import java.util.Collections;
import java.util.PriorityQueue;

class MedianFinder {
    private PriorityQueue<Integer> left;
    private PriorityQueue<Integer> right;

    public MedianFinder() {
        left = new PriorityQueue<>(Collections.reverseOrder());
        right = new PriorityQueue<>();
    }

    public void addNum(int num) {
        if (left.isEmpty() || num <= left.peek()) {
            left.offer(num);
        } else {
            right.offer(num);
        }

        if (left.size() > right.size() + 1) {
            right.offer(left.poll());
        } else if (right.size() > left.size()) {
            left.offer(right.poll());
        }
    }

    public double findMedian() {
        if (left.size() > right.size()) {
            return left.peek();
        }
        return (left.peek() + right.peek()) / 2.0;
    }
}

为什么两个堆能维护中位数

中位数只关心：

text 复制代码

一半更小的数
一半更大的数

两个堆刚好分别维护这两半，并且堆顶能快速取到边界值。

双堆把数据一分为二，左边用大根堆，右边用小根堆。

题型六：滑动窗口最大值

这是堆的一个经典进阶应用。

题目：

给定一个数组和窗口大小 k，返回每个窗口中的最大值。

为什么能想到堆

窗口在移动时，最需要维护的是：

text 复制代码

当前窗口里的最大值

所以可以用大根堆存下标和数值，在每次窗口右移时快速取堆顶。

Java 代码实现

java 复制代码

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

class Solution {
    public int[] maxSlidingWindow(int[] nums, int k) {
        PriorityQueue<int[]> pq = new PriorityQueue<>((a, b) -> b[0] - a[0]);
        List<Integer> ans = new ArrayList<>();

        for (int i = 0; i < nums.length; i++) {
            pq.offer(new int[]{nums[i], i});

            while (!pq.isEmpty() && pq.peek()[1] <= i - k) {
                pq.poll();
            }

            if (i >= k - 1) {
                ans.add(pq.peek()[0]);
            }
        }

        int[] res = new int[ans.size()];
        for (int i = 0; i < ans.size(); i++) {
            res[i] = ans.get(i);
        }
        return res;
    }
}

这道题的关键点

堆里存的不只是数值，还要存下标。

因为窗口滑动后，过期元素必须被剔除。

如果只存数值，你无法判断堆顶是不是已经不在当前窗口里了。

堆题里常要同时保存数值和下标，方便判断元素是否过期。

题型七：数据流中的第 K 大元素

这类题非常贴近"在线维护"的实际需求。

题目：

设计一个数据结构，支持不断加入数字，并返回当前第 k 大元素。

核心思路

维护一个大小为 k 的小根堆：

堆里保存当前最大的 k 个元素
堆顶就是这 k 个元素里最小的，也就是第 k 大元素

Java 代码实现

java 复制代码

import java.util.PriorityQueue;

class KthLargest {
    private final int k;
    private final PriorityQueue<Integer> pq;

    public KthLargest(int k, int[] nums) {
        this.k = k;
        this.pq = new PriorityQueue<>();

        for (int num : nums) {
            add(num);
        }
    }

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

        return pq.peek();
    }
}

为什么这题特别适合堆

因为它是"数据流"场景。

数据不断进来时，堆可以持续维护一个固定大小的最优集合。

这比每次重新排序更高效，也更符合在线问题的处理方式。

固定大小的小根堆非常适合动态维护最大的 K 个元素。

进阶视角：堆和排序不是替代关系

堆和排序经常一起出现，但它们解决的问题不同。

排序适合什么

排序适合你要：

一次性得到全局顺序
所有元素都要处理

堆适合什么

堆适合你要：

只关心最值
只保留前 K 个
数据是流式进入的
需要不断插入和弹出

一个直观对比

如果你每次都做全排序：

text 复制代码

成本高，但结果完整

如果你用堆：

text 复制代码

只维护最关键部分，结果更聚焦

排序是全局整理，堆是局部维护。

常见坑点：堆题最容易错在哪里

1. 方向用反

很多 TopK 题里最容易错的是：

该用小根堆却写成大根堆
该用大根堆却写成小根堆

记住一个简单原则：

text 复制代码

保留最大的 K 个，用小根堆
保留最小的 K 个，用大根堆

2. 忘记堆顶代表什么

堆顶不是"排名第一"的绝对含义，而是：

小根堆：当前堆内最小值
大根堆：当前堆内最大值

题目想保留哪一侧的边界，就选哪种堆。

3. 比较器写错

比如频率堆应该按频率排序，不是按元素值排序。

4. Java `PriorityQueue` 默认不是大根堆

默认是小根堆。

如果题目需要大根堆，要显式传入：

java 复制代码

Collections.reverseOrder()

5. 堆里存什么没想清楚

有些题只存数值就够了。

有些题需要存：

数值和频率
数值和下标
节点对象

先想清楚"堆节点里到底需要保存哪些信息"，再写代码。

复杂度分析：堆为什么适合 TopK

假设有 n 个元素，想找 TopK。

直接排序

text 复制代码

O(n log n)

用堆维护 K 个元素

每个元素最多做一次入堆和一次出堆：

text 复制代码

O(n log k)

当 k << n 时，O(n log k) 通常明显优于 O(n log n)。

空间复杂度

TopK：O(k)
频率统计：O(m)，m 是不同元素个数
双堆中位数：O(n)，取决于数据规模和维护方式

堆不会减少全部数据的读取成本，但能把"只关心前 K 个"的代价压到 O(n log k)。

总结

堆不是用来做全量排序的。

它最擅长的是在数据持续进入时，快速保留最有价值的那一小部分。

你可以重点记住下面几句话：

堆只保证堆顶是当前最值
Java 的 PriorityQueue 默认是小根堆
前 K 大用小根堆
前 K 小用大根堆
频率 TopK 先统计再进堆
合并 K 个有序序列时，堆负责拿当前最小头部
动态中位数常用双堆维护
堆题的核心是先想清楚"堆里该存什么"
堆的价值在于把全量比较变成局部维护

第23篇-堆与优先队列-TopK问题的常规武器

概述

核心概念：堆到底是什么

小根堆示意

大根堆示意

堆和排序的区别

优先队列：Java 里怎么用堆

小根堆写法

大根堆写法

常用 API

为什么 PriorityQueue 适合刷题

题型一：数组中的前 K 大元素

解题思路

Java 代码实现

为什么最后要倒序取出

题型二：数组中的前 K 小元素

解题思路

Java 代码实现

为什么前 K 小用大根堆

题型三：出现频率最高的 K 个元素

解题思路

Java 代码实现

为什么堆里存的是数组

复杂度分析

题型四：合并 K 个有序数组或链表

合并 K 个升序链表示例

为什么每个链表只放头节点

题型五：动态维护中位数

核心思想

代码示意

为什么两个堆能维护中位数

题型六：滑动窗口最大值

为什么能想到堆

Java 代码实现

这道题的关键点

题型七：数据流中的第 K 大元素

核心思路

Java 代码实现

为什么这题特别适合堆

进阶视角：堆和排序不是替代关系

排序适合什么

堆适合什么

一个直观对比

常见坑点：堆题最容易错在哪里

1. 方向用反

2. 忘记堆顶代表什么

3. 比较器写错

4. Java PriorityQueue 默认不是大根堆

5. 堆里存什么没想清楚

复杂度分析：堆为什么适合 TopK

直接排序

用堆维护 K 个元素

空间复杂度

总结

为什么 `PriorityQueue` 适合刷题

4. Java `PriorityQueue` 默认不是大根堆