Lesson 0406 · 堆与优先队列

④ 数据结构与算法 — 第 6 课

1. 从你写过的题说起

P1090 合并果子：有 n 堆果子，每次把两堆合并成一堆，消耗体力 = 两堆重量之和，求合并成一堆的最小总体力。

贪心策略很直白：每次挑最轻的两堆合并。于是每一步都要做两件事——

• 找出当前的最小值（还要找出第二小）
• 删掉这两个，塞进去一个新的重量

如果用 sort 每轮重排一次：一轮 O(n log n)，总共 O(n² log n)，n 上万就超时。如果线性扫一遍找最小：一轮 O(n)，总共 O(n²)，也不够快。

堆就是为「动态地反复取极值」这件事设计的。它比 BST 简单得多，因为它只维护一小部分信息。

2. 堆是什么：完全二叉树 + 一条规矩

堆（heap）满足两个条件：

• 形状：是一棵完全二叉树——除了最后一层，每层都填满，最后一层从左往右连续填。
• 规矩：每个节点的值都 ≤ 它的两个孩子（这叫小根堆，堆顶是最小值）；反过来每个节点都 ≥ 孩子，就是大根堆。

小根堆示例：              对应的数组（1-based）：
        1                 h[1]=1
       / \                h[2]=3  h[3]=2
      3   2               h[4]=7  h[5]=5  h[6]=4
     / \  /
    7   5 4

注意两件事，这是初学者最容易绕不过去的弯：

1. 堆顶一定是最小值，但堆里其他位置不是有序的。上面数组写成 1 3 2 7 5 4，看着乱，但 1 永远在顶上。
2. 堆只保证父子之间的大小关系，兄弟之间没有任何约束。3 和 2 谁在左谁在右都行。

所以堆和 BST 是两种不同的取舍：

3. 用数组装堆

完全二叉树的形状规则，意味着它没有空洞，可以直接用数组紧凑地存，不需要指针。

约定从下标 1 开始存（根放在 h[1]），下标关系就非常干净：

节点 i 的左孩子： 2*i
节点 i 的右孩子： 2*i + 1
节点 i 的父亲：   i / 2      （整数除法，自动向下取整）

下标：  1   2   3   4   5   6
值：    1   3   2   7   5   4

h[2]=3 的孩子是 h[4]=7、h[5]=5
h[3]=2 的孩子是 h[6]=4（h[7] 不存在，说明没有右孩子）

4. 两个动作就够了

4.1 插入：先放到末尾，再「上浮」

新元素先放在数组末尾（这样仍然是完全二叉树），然后不停和父亲比：比父亲小，就和父亲交换，直到不比父亲小为止。

const int N = 100005;
int h[N], hn = 0;                  // hn = 当前元素个数，堆顶是 h[1]

void push(int x) {
    h[++hn] = x;                   // 先放到末尾
    int i = hn;
    while (i > 1 && h[i] < h[i / 2]) {   // 比父亲小就往上换
        swap(h[i], h[i / 2]);
        i /= 2;                    // 继续和新的父亲比
    }
}

把原文那个例子逐轮拆开（往 1 3 2 7 5 4 里插 0）：

三轮之后 0 浮到了堆顶。上浮只走一条从叶到根的路径，长度是树高，所以是 O(log n)。

4.2 删除堆顶：把末尾搬上来，再「下沉」

删掉堆顶之后，让最后一个元素搬到堆顶（还是完全二叉树），然后不停和自己较小的那个孩子比：比孩子大就交换，直到不比孩子大。

int top() { return h[1]; }         // 看最小值：O(1)

void pop() {
    h[1] = h[hn--];                // 末尾元素搬到堆顶，个数减一
    int i = 1;
    while (true) {
        int l = i * 2, r = i * 2 + 1;
        int small = i;                     // 假设自己最小
        if (l <= hn && h[l] < h[small]) small = l;
        if (r <= hn && h[r] < h[small]) small = r;
        if (small == i) break;             // 两个孩子都不比我小，停
        swap(h[i], h[small]);
        i = small;                         // 继续往下沉
    }
}

还是用上面那个堆 1 3 2 7 5 4，删一次堆顶：

一次上浮或下沉，最多走一条从根到叶的路径，长度是树高 log n，所以两种操作都是 O(log n)。top() 是 O(1)。

5. STL 里现成的优先队列

堆在 STL 里的名字叫 priority_queue（优先队列），头文件是 <queue>。

#include <queue>
using namespace std;

priority_queue<int> big;                                      // 大根堆（默认）
priority_queue<int, vector<int>, greater<int>> small;        // 小根堆

small.push(3);
small.push(1);
small.push(5);

cout << small.top() << endl;    // 1，堆顶是最小值
small.pop();                     // 删掉 1（注意 pop 不返回值）
cout << small.top() << endl;    // 3
cout << small.size() << endl;   // 1
cout << small.empty() << endl;  // 0

5.1 自定义类型：按分数取最高

struct Student {
    string name;
    int score;
};

struct Cmp {                              // 自定义「谁优先级低」
    bool operator()(const Student& a, const Student& b) {
        return a.score < b.score;
    }
};

priority_queue<Student, vector<Student>, Cmp> pq;
pq.push({"Alice", 92});
pq.push({"Bob",   78});
pq.push({"Cara",  99});

cout << pq.top().name;     // Cara（分数最高在堆顶）

更省事的写法是在结构体里直接重载 <：

struct Student {
    string name;
    int score;
    bool operator<(const Student& o) const {
        return score < o.score;           // 分数大的排前面
    }
};

priority_queue<Student> pq;               // 不用再写第三个参数

6. 什么时候用堆

7. 两个典型用法

7.1 合并果子（P1090 完整代码）

#include <iostream>
#include <queue>
#include <vector>
using namespace std;

int main() {
    int n;
    cin >> n;

    priority_queue<int, vector<int>, greater<int>> q;   // 小根堆
    for (int i = 0; i < n; i++) {
        int w;
        cin >> w;
        q.push(w);
    }

    long long ans = 0;                       // 体力总和可能超过 int
    while (q.size() > 1) {
        int a = q.top(); q.pop();            // 最轻的一堆
        int b = q.top(); q.pop();            // 第二轻的
        ans += a + b;
        q.push(a + b);                       // 合并后放回去
    }

    cout << ans << endl;
    return 0;
}

样例 3 / 1 2 9 的输出是 15（先合并 1+2=3 花 3 点体力，再合并 3+9=12 花 12 点，共 15）。

7.2 只要最大的 k 个

思路反着来：用一个小根堆只保留 k 个数，多出来的就把最小的扔出去，剩下的自然就是最大的 k 个。

int k = 3;
priority_queue<int, vector<int>, greater<int>> keep;

for (int x : v) {
    keep.push(x);
    if ((int)keep.size() > k) keep.pop();    // 超过 k 个就扔掉最小的
}

cout << keep.top() << endl;   // 堆顶是第 k 大的数

反过来求「最小的 k 个」就用大根堆，其它一字不改。

8. 坑点清单

• 下标基准混用：2*i、i/2 是 1-based 的公式；从 0 开始要换成 2i+1、2i+2、(i-1)/2，别只改一半。
• 空堆还 pop()：h[1] = h[hn--] 在 hn = 0 时读到的是 h[0] 的垃圾值，取堆顶前先判空。
• 忘记更新 hn：push 后不加、pop 后不减，后面所有边界判断全错位。
• 下沉时跟「较大的孩子」换：必须跟较小的那个换，否则换完父亲还是比另一个孩子大，堆序被破坏。
• 循环边界写反：上浮是 i > 1，下沉是 l <= hn；差一个等号就会越界或提前停。
• STL 的两个方向坑：pop() 不返回值；自定义比较器和 sort 的 cmp 方向正好相反。

9. 小测

10. 动手练习

11. 本课小结

• 堆 = 完全二叉树 + 父子顺序规矩；小根堆堆顶最小，大根堆堆顶最大。
• 堆只保证极值在顶上，遍历出来不是有序的——这是它比 BST 简单的原因。
• 用 1-based 数组存：左孩子 2i、右孩子 2i+1、父亲 i/2。
• 插入靠上浮，删堆顶靠下沉，都是 O(log n)；看堆顶 O(1)。
• STL 写法：priority_queue<int> 大根堆，priority_queue<int, vector<int>, greater<int>> 小根堆。
• 自定义比较的方向和 sort 的 cmp 相反，容易搞反。
• 典型场景：反复取极值、合并果子这类贪心、Top-K、对顶堆求中位数。