跳转至

第 35 章 优先队列与堆

配套例题:BISHI6 【模板】整数优先队列、BISHI49 小红闯关、BISHI50 [JSOI2007] 建筑抢修、BISHI51 低买高卖(题解均已通过官方样例) 来源:S2 useful algorithm/heapsort.cpp 前置33-队列与双端队列12-自定义排序

优先队列不是「排好序的队列」,而是只保证能 \(O(1)\) 看到最值、\(O(\log n)\) 取出最值的容器。 Python 的 heapq 是 C 实现的二叉堆,是标准库里为数不多的竞赛级快的数据结构。


35.1 二叉堆的原理

完全二叉树 + 堆序性质:每个节点都不大于(小根堆)其两个孩子。

用数组存完全二叉树,不需要指针。0-indexed 时:

关系 下标
节点 \(i\) 的左孩子 \(2i+1\)
节点 \(i\) 的右孩子 \(2i+2\)
节点 \(i\) 的父亲 \(\lfloor (i-1)/2 \rfloor\)

S2 的 heapsort.cpp 用的是 1-indexed(左孩子 \(2i\)、右孩子 \(2i+1\)、父亲 \(i/2\)), 公式更漂亮,代价是浪费下标 0。heapq 用的是 0-indexed。

两个基本操作

操作 别名 做法 复杂度
上浮 sift up / 向上调整 新元素放到末尾,不断和父亲比较交换 \(O(\log n)\)
下沉 sift down / 向下调整 根被替换后,不断和较小的孩子交换 \(O(\log n)\)

S2 的 head(max, num) 函数就是下沉(那份代码建的是大根堆):

j = 2 * num;                              // 左孩子
while(j <= max){
    if(j < max && (a[j] < a[j+1])) j++;   // 取较大的孩子
    if(a[num] < a[j]){                    // 父亲比孩子小,交换后继续下沉
        swap(a[num], a[j]);
        num = j; j = 2 * num;
    }
    else j = n + 1;                       // 已满足堆序,退出
}

易错点:下沉时必须和两个孩子中更优的那个比较, 只和左孩子比会破坏堆序。if(j < max && a[j] < a[j+1]) j++ 这一行就是在选孩子。

heapify 为什么是 \(O(n)\)

从最后一个非叶节点开始,倒着对每个节点做一次下沉:

for (i = n / 2; i >= 1; i--) head(n, i);

看起来是 \(n/2\)\(O(\log n)\) = \(O(n\log n)\),但深度越大的节点越多、下沉距离越短

\[\sum_{h=0}^{\log n} \frac{n}{2^{h+1}} \cdot h = n \sum_{h=0}^{\log n} \frac{h}{2^{h+1}} < n \cdot 1 = O(n)\]

所以 heapq.heapify(a)\(O(n)\),比逐个 heappush\(O(n\log n)\) 快。 已知全部元素时永远用 heapify


35.2 heapq 全 API

Python 的 heapq 是小根堆h[0] 永远是最小值),且直接作用在普通 list

import heapq

h = []
heapq.heappush(h, x)                # 插入        O(log n)
x = heapq.heappop(h)                # 弹出最小值   O(log n)
x = h[0]                            # 看最小值     O(1),不删除
heapq.heapify(a)                    # 就地建堆     O(n)
len(h)                              # 元素个数     O(1)

heapq.heappushpop(h, x)             # 先 push 再 pop,O(log n),比分开写快
heapq.heapreplace(h, x)             # 先 pop 再 push,O(log n)
heapq.nlargest(k, a)                # 前 k 大,O(n log k)
heapq.nsmallest(k, a)               # 前 k 小,O(n log k)
heapq.merge(*iterables)             # 归并多个有序序列,返回迭代器,O(总长 log k)

heappushpop vs heapreplace

函数 语义 何时用
heappushpop(h, x) 先插入 \(x\),再弹最小 「维护前 \(k\) 大」的标准写法
heapreplace(h, x) 先弹最小,再插入 \(x\) 已知 \(x\) 一定大于堆顶时,省一次比较

两者都比 heappush + heappop 快约 30%(只做一次下沉,不做上浮)。

# 维护「前 k 大」的 O(n log k) 标准模板
h = []                               # 小根堆:堆顶是这 k 个候选里最小的那个
for x in a:
    if len(h) < k:                   # 还没攒够 k 个,来者不拒
        heapq.heappush(h, x)
    elif x > h[0]:                   # 只有比现任第 k 大更大,才值得挤进来
        heapq.heapreplace(h, x)      # 堆里只留 k 个最大的,堆顶是第 k 大
        # x <= h[0] 时直接丢弃:它连当前的第 k 名都挤不掉,更不可能进最终答案

注意 heapq 不提供的东西: - ❌ 没有大根堆 - ❌ 没有「删除任意元素」 - ❌ 没有「修改某个元素的值」(decrease-key) - ❌ 没有「合并两个堆」(h1 + h2 后要重新 heapify

这四件事各有替代方案,见下文。


35.3 最大堆的三种做法

做法 写法 评价
取负 heappush(h, -x),取出时 -heappop(h) ✅ 首选,最快
heapq._heapify_max 等私有函数 heapq._heapify_max(a) ⚠️ 私有 API,只有 _heapify_max/_heappop_max没有 _heappush_max,不好用
包装类实现 __lt__ 自定义反序比较 ❌ 每次比较都要调 Python 方法,慢 5 倍
import heapq

# 大根堆:全部取负
h = []
heapq.heappush(h, -x)
mx = -h[0]                          # 看最大值
mx = -heapq.heappop(h)              # 取最大值

取负的两个坑: 1. 元组也要逐项取负(-w, u) 而不是 -(w, u)(后者直接 TypeError); 2. 只取负第一维时,第二维的比较方向不变—— (-w, u)\(w\) 相同时按 \(u\) 升序,不是降序。要全反就写 (-w, -u)


35.4 堆里放什么:元组、类、还是 key

heapq 没有 key 参数(和 sorted 不同)。要按自定义规则排序,只有两条路。

路线一:放元组(首选

Python 的元组比较是逐元素字典序,天然支持多关键字:

heapq.heappush(h, (dist, node))              # 按 dist 升序,dist 相同按 node 升序
heapq.heappush(h, (-profit, cost, idx))      # profit 降序,cost 升序,idx 升序

陷阱:元组比较到第 \(k\) 项时,如果前 \(k-1\) 项全相等,就会比较第 \(k\) 项。 如果第 \(k\) 项是不可比较的对象(比如自定义类),会抛 TypeError

heapq.heappush(h, (d, obj))     # ❌ d 相同时会去比较两个 obj
heapq.heappush(h, (d, i, obj))  # ✅ 加一个唯一的整数 i 当「打破平局键」

这是「堆里放对象」时最常见的运行时错误。

路线二:给类实现 __lt__

import functools


@functools.total_ordering            # 由 __lt__ 和 __eq__ 自动补出 >、<=、>= 等
class Task:
    __slots__ = ("d", "t", "idx")    # 固定属性表:省掉实例字典,取属性也更快

    def __init__(self, d, t, idx):
        self.d = d
        self.t = t
        self.idx = idx               # 唯一编号,用于在 d 与 t 都相同时打破平局

    def __lt__(self, other):
        # heapq 只用 < 一个运算符,所以只实现 __lt__ 就够
        if self.d != other.d:        # 第一关键字:d 升序
            return self.d < other.d
        return self.t < other.t      # d 相同时看第二关键字 t

能用元组就别用类:类的 __lt__ 是 Python 层函数调用, \(10^5\)heappush 会产生约 \(1.7\times10^6\) 次调用,比元组比较慢 5 倍以上。 详见 15-面向对象与迭代器生成器


35.5 堆排序

import heapq


def heap_sort(a):
    """堆排序,O(n log n),不稳定。教学用途 —— 实战永远用 sorted()。"""
    h = list(a)
    heapq.heapify(h)                          # O(n)
    return [heapq.heappop(h) for _ in range(len(h))]   # n 次 O(log n)
排序 复杂度 稳定 额外空间 Python 实测
sorted()(Timsort) \(O(n\log n)\) \(O(n)\) 基准 1×
heap_sort(上面) \(O(n\log n)\) \(O(n)\) 约 8× 慢
手写堆排(纯 Python 上浮下沉) \(O(n\log n)\) \(O(1)\) 约 60× 慢

结论:在 Python 里手写任何排序算法都只有教学意义。 实战一律 sorted()。见 40-排序

堆真正不可替代的场景是「动态」:元素边进边出,随时要最值。 静态数据全部就位时,排序永远比堆好。


35.6 堆的四大套路

套路 识别特征 模板位置
\(k\) 大 / 第 \(k\) 「最大的 \(k\) 个」 35.2
反悔贪心 先全都要,超限了再退掉最差的 35.7(BISHI50/51)
懒删除堆 需要删除任意元素 34-集合与多重集合 34.7
对顶堆 动态维护中位数 本节
堆优化 Dijkstra 带权最短路 91-最短路
多路归并 合并 \(k\) 个有序序列 heapq.merge

对顶堆:动态维护中位数

用一个大根堆存较小的一半,一个小根堆存较大的一半, 保持两堆大小差不超过 1,中位数就在堆顶。

import heapq


class MedianKeeper:
    """动态维护中位数。add 均摊 O(log n),median O(1)。"""

    def __init__(self):
        self.lo = []            # 大根堆(存负数),保存较小的一半
        self.hi = []            # 小根堆,保存较大的一半

    def add(self, x):
        # -lo[0] 是较小一半里的最大值,即当前的分界线;不超过它就归 lo
        if self.lo and x < -self.lo[0]:
            heapq.heappush(self.lo, -x)      # 存进大根堆要取负
        else:
            heapq.heappush(self.hi, x)
        # 再平衡:|lo| 与 |hi| 相差不超过 1,且 |lo| >= |hi|
        if len(self.lo) > len(self.hi) + 1:
            heapq.heappush(self.hi, -heapq.heappop(self.lo))    # 弹出时再取负还原
        elif len(self.hi) > len(self.lo):
            heapq.heappush(self.lo, -heapq.heappop(self.hi))
        # 每次至多搬一个元素:单次插入只会让某边多出 1,不会破坏更多

    def median(self):
        if len(self.lo) > len(self.hi):      # 总数为奇数时中位数落在 lo 顶上
            return -self.lo[0]
        return (-self.lo[0] + self.hi[0]) / 2        # 偶数则取两个堆顶的平均

反悔贪心:本章最重要的套路

模式是固定的:

先无脑把每个元素都收下,一旦违反约束,就从已收的里面弹掉「最亏」的那个。

它之所以正确,是因为「弹掉一个换进一个」的操作不会让已选个数变少, 所以每一步都保持着「在当前前缀下能选到的最优解」。这是标准的交换论证, 详见 47-贪心


35.7 例题

BISHI6 【模板】整数优先队列(简单)

\(n \le 10^6\) 次操作:1 x 插入;2 输出当前最小元素;3 删除一个最小元素。 题面见 BISHI6 原题(牛客)

heapq 本来就是小根堆,不用取负,直接对应。 这题的真正考点是 \(n \le 10^6\) 下的常数优化:算法是 \(O(n\log n)\) 没得挑, 瓶颈全在 IO 和解释器开销。

三个优化点:

  1. 一次性读入全部 token 用游标推进,绝不用 \(10^6\)input()
  2. heappush/heappop 绑成局部变量——\(10^6\) 次调用能省下可观的属性查找开销 (LOAD_FAST vs LOAD_GLOBAL + LOAD_ATTR);
  3. 输出攒进列表最后一次性 write
import sys
from heapq import heappush, heappop


def main():
    data = sys.stdin.buffer.read().split()   # 1e6 次 input() 会拖垮整个程序
    n = int(data[0])
    i = 1                                # token 游标:操作长度不固定,只能按 token 走
    h = []                               # heapq 直接作用在普通 list 上,无需包装
    out = []
    push, pop = heappush, heappop        # 绑成局部名
    for _ in range(n):
        op = data[i]                     # data 未解码,比较对象是 bytes
        i += 1
        if op == b"1":                   # 插入:本条操作占 2 个 token
            push(h, int(data[i]))
            i += 1                       # 多吃掉一个操作数
        elif op == b"2":                 # 查询最小值,不删
            if h:
                out.append(h[0])         # 直接读堆顶,pop 再 push 要多两次 O(log n)
        else:                            # 3:删除一个最小值,不输出
            if h:
                pop(h)                   # 判空只为防 RE,题面未保证操作时堆非空
    sys.stdout.write("\n".join(map(str, out)) + "\n")    # 一次性输出,省下 1e6 次系统调用


main()

四个坑

  1. 操作 2/3 只有一个 token,操作 1 有两个,行长度不固定,必须用游标按 token 读;
  2. 操作 2 是查询(不删),操作 3 是删除(不输出),只有操作 2 产生输出;
  3. 查询最小值直接看 h[0],不要 heappop 之后再 heappush—— 那要多两次 \(O(\log n)\) 调整,白白慢一倍;
  4. 空堆时的 2/3 题面没说会出现,但加个判空更稳,免得 RE。

题解见 solutions/BISHI6.py

BISHI49 小红闯关(中等)

\(n\) 个关卡按顺序通过,第 \(i\) 关耗时 \(a_i\)每通过 \(k\) 个关卡获得一个跳关道具 (跳关也算一次成功闯关),道具可在任意关卡使用,使用后不消耗时间直接通过。 求通过全部关卡的最少总时间。\(1 \le n, k \le 10^5\)\(1 \le a_i \le 10^5\)。 题面见 BISHI49 原题(牛客)

建模:设跳过的关卡集合为 \(S\),答案 \(= \sum a_i - \sum_{i \in S} a_i\), 所以要最大化被跳过关卡的耗时之和

约束是什么? 要在第 \(i\) 关用道具,必须在此之前已经攒到一个没用掉的。 通过前 \(i-1\) 关后共获得 \(\lfloor (i-1)/k \rfloor\) 个道具,于是对任意前缀 \(i\)

\[|S \cap [1, i]| \le \left\lfloor \frac{i-1}{k} \right\rfloor\]

这是一个前缀容量约束,容量随 \(i\) 单调不减——正是反悔贪心的标准形态。

import sys
from heapq import heappush, heappushpop


def main():
    data = sys.stdin.buffer.read().split()
    n = int(data[0]); k = int(data[1])
    a = [int(v) for v in data[2:2 + n]]

    heap = []                       # 小根堆:当前决定跳过的关卡耗时
    total = 0                       # 全部关卡耗时之和,最后减去被跳过的部分
    for i in range(1, n + 1):       # i 用 1 起,与前缀容量公式保持一致
        t = a[i - 1]                # a 是 0 下标,故取 i-1
        total += t
        cap = (i - 1) // k          # 进入第 i 关之前手上的道具总数
        if len(heap) < cap:         # 还有空余道具,直接收下
            heappush(heap, t)
        elif cap and t > heap[0]:   # 道具已用满:只有比最小的那个更值才换
            heappushpop(heap, t)    # 容量满了,换掉最小的那个
            # cap 为 0 时一个都跳不了,这一项必须挡住,否则会凭空多跳一关
    print(total - sum(heap))        # 堆里留下的就是最终跳过的集合


main()

复杂度 \(O(n \log n)\)\(n = 10^5\) 时约 \(1.7\times10^6\) 次堆调整,在 2 秒限制下没问题。

手动验证样例 2n=6, k=2, a = 1 1 4 5 1 4total = 16):

\(i\) \(\text{cap} = \lfloor (i-1)/2 \rfloor\) \(a_i\) 动作 saved
1 0 1 满(0 个位置) [] 0
2 0 1 [] 0
3 1 4 有空位,收 [4] 4
4 1 5 满,\(5 > 4\),换 [5] 5
5 2 1 有空位,收 [1,5] 6
6 2 4 满,\(4 > 1\),换 [4,5] 9

答案 \(16 - 9 = 7\) ✓(题目样例输出 7)

四个坑

  1. 跳关本身也算「通过一关」,所以计数照常推进。 样例 3(n=5, k=1, a = 2 4 5 1 3)就是在考这个:\(k=1\) 时打完第一关, 后面全能跳,答案只有 \(a_1 = 2\)
  2. 容量是 \(\lfloor (i-1)/k \rfloor\) 而不是 \(\lfloor i/k \rfloor\)—— 进第 \(i\) 关时第 \(i\) 关还没通过。用后者会多算一个道具,样例 1 会错成 1;
  3. 道具可以攒着不用、也可以在任意后续关卡使用,只有前缀数量约束, 没有「必须立刻用掉」的限制;
  4. \(n = k\)\(c_n = \lfloor (n-1)/k \rfloor = 0\),一个都跳不了。

为什么不能「先把所有 \(a_i\) 塞进大根堆,容量增加时弹最大的」? 因为那样会追溯性地跳过前面的关卡——比如样例 3 里会跳过第 1 关, 但那时候一个道具都还没有。贪心必须尊重时间顺序, 所以要从左往右扫、用小根堆做「已选集合」的反悔。

题解见 solutions/BISHI49.py

BISHI50 [JSOI2007] 建筑抢修(中等)

\(n \le 1.5\times10^5\) 座建筑,修第 \(i\) 座要 \(t_i\) 秒,必须在第 \(d_i\)之前修完, 否则报废。一次只能修一座(可瞬移)。求最多能修好几座。 题面见 BISHI50 原题(牛客)

教科书级的反悔贪心,和 BISHI49 是同一个模子:

  1. 按截止时间 \(d_i\) 升序排序;
  2. 依次尝试修每一座,累计用时 cur += t_i,并把 \(t_i\) 推入大根堆
  3. 一旦 cur > d_i(来不及了),就从堆里弹出耗时最长的那座放弃,cur -= 弹出值
  4. 答案 = 最终堆的大小。

正确性:放弃耗时最长的那座,能腾出最多时间且只损失一座, 而已选个数不减,所以任何时刻堆的大小都是「前 \(i\) 座里能修的最大数量」。

为什么按 \(d\) 排序? 若某个可行解里存在 \(d\) 大的排在 \(d\) 小的前面, 交换这两个任务不会让任何一个超时(经典交换论证), 所以总存在一个按 \(d\) 升序执行的最优解。

import sys
from heapq import heappush, heappop


def main():
    data = sys.stdin.buffer.read().split()
    n = int(data[0])
    jobs = []
    p = 1                             # token 游标:每座建筑两个数
    for _ in range(n):
        t = int(data[p]); d = int(data[p + 1]); p += 2
        jobs.append((d, t))           # 存成 (d, t):排序键必须是 d,放第一位
    jobs.sort()                       # 按截止时间升序

    heap = []                         # 大根堆(存负数):已接下的任务耗时
    cur = 0                           # 已接任务的总耗时,也就是当前时刻
    for d, t in jobs:
        cur += t                      # 先无条件接下,违约了再反悔
        heappush(heap, -t)
        if cur > d:                   # 修不完了,反悔掉最耗时的那个
            cur += heappop(heap)      # heap 里是负数,加上等于减去耗时
            # 取等号(cur == d)算按时完成,所以判据是严格大于
            # 弹出的可能正是刚推进去的这一座,无需回退循环
    print(len(heap))                  # 堆的大小恒等于已选任务数,即答案


main()

复杂度 \(O(n \log n)\)

手动验证样例(按 \(d\) 排序后是 \((100,200), (1000,1250), (200,1300), (2000,3200)\), 这里写成 \((t, d)\)):

顺序 \(t\) \(d\) cur 判定
1 100 200 100 \(100 \le 200\)
2 1000 1250 1100 \(1100 \le 1250\)
3 200 1300 1300 \(1300 \le 1300\)取等号也算成功
4 2000 3200 3300 \(3300 > 3200\),弹出最大的 2000,cur 回到 1300

堆里剩 3 个 → 答案 3 ✓

四个坑

  1. 排序键只用 \(d\),不能用 \(t\)\(d - t\)
  2. 判定是 cur > d 才弹,也就是 cur <= d 算成功(取等号也算)。 样例里恰好有一个 \(1300 = 1300\) 的边界,写成 cur >= d 就会少算一个;
  3. 只在超时时弹,且弹完不用回退指针——当前任务已经在堆里, 弹出的可能正是它自己(说明它太长,不如不修);
  4. \(d_i \le 2\times10^9\) 超过 32 位——C++ 要开 long longPython 无所谓, 这是 Python 少见的优势场景。

题解见 solutions/BISHI50.py

BISHI51 低买高卖(中等)

已知 \(n \le 3\times10^5\) 天的股价 \(p_i\)。每天至多操作一次:买入 1 股、卖出 1 股、或不动。 初始无股票、不能做空,第 \(n\) 天收盘必须清空。求最大收益。 题面见 BISHI51 原题(牛客)

这题的反悔贪心最精巧,值得单独理解。

import heapq
import sys


def main():
    data = sys.stdin.buffer.read().split()
    n = int(data[0])
    ans = 0
    h = []                      # 小根堆:所有可用的买入价(含反悔票)
    push, pop = heapq.heappush, heapq.heappop        # 3e5 天要调 6e5 次,绑局部名
    for tok in data[1:n + 1]:   # 只取前 n 个价格,尾部多余 token 不参与
        p = int(tok)
        if h and h[0] < p:      # 存在比今天更低的买入价,这笔差价就值得先吃下
            ans += p - pop(h)   # 以堆顶价买入、今天 p 卖出,先吃下这笔
            push(h, p)          # ★ 反悔票:允许把这笔交易的卖出价往后挪
        push(h, p)              # p 自身也可以充当未来的买入价
        # 两次 push 缺一不可:少了反悔票就退化成「全程只做一笔交易」的错解
    sys.stdout.write(str(ans) + "\n")   # 收益可达 1.5e11,Python 整数无溢出之忧


main()

那个「多推一次 p」是全题的灵魂。

假设某天以 \(p\) 卖掉了成本 \(q\) 的股票,赚 \(p - q\)。 但也许后面有个更高的价 \(r > p\),真正的最优是 \(q\) 买、\(r\) 卖。 注意到

\[(p - q) + (r - p) = r - q\]

也就是说:只要把 \(p\) 重新放回堆里当作「一个成本为 \(p\) 的买入机会」, 将来以 \(r\) 卖出时补上的 \(r - p\) 就恰好把答案修正成 \(r - q\)。 「反悔」不需要真的撤销操作,只要让差价能拼接起来。

另一个 heappush(h, p) 是常规的「今天可以买入」。所以每天最多推两次。

手动验证样例p = 10 5 4 7 9 12 6 2 10,答案 20):

\(p\) 堆顶 动作 ans
1 10 只推入 0
2 5 10 \(5 < 10\),不卖 0
3 4 5 不卖 0
4 7 4 卖:\(+3\),推回 7 3
5 9 5 卖:\(+4\),推回 9 7
6 12 7 卖:\(+5\),推回 12 12
7 6 7 \(6 < 7\),不卖 12
8 2 6 不卖 12
9 10 2 卖:\(+8\) 20

答案 20 ✓(题解给的方案是 \((9-5)+(12-4)+(10-2) = 20\),数值一致)

为什么「第 \(n\) 天必须清空」不需要额外处理? 因为算法从不「持有未平仓的股票」——每一笔加进 ans 的都是已完成的差价对, 堆里剩下的只是没被使用的买入候选。自动满足约束。

三个坑

  1. 必须压两次(一次「反悔票」+ 一次「自身作为买点」), 只压一次就退化成「只能做一笔」的错解;
  2. 「第 \(n\) 天必须清仓」不需要额外处理,见上;
  3. 收益最大可达 \(\frac{3\times10^5}{2} \times 10^6 \approx 1.5\times10^{11}\), C++ 要 long long,Python 无所谓。

复杂度 \(O(n \log n)\)\(3\times10^5\) 天最多 \(6\times10^5\)heappushheapq 底层是 C 实现,毫秒级。

题解见 solutions/BISHI51.py


35.8 本章速查

要点 结论
heapq小根堆 h[0] 是最小值
大根堆 全部取负,元组要逐项取负
heapify \(O(n)\),不是 \(O(n\log n)\)
看最值 h[0]\(O(1)\)不要 pop 再 push
heappushpop / heapreplace 比分开写快约 30%
\(k\) 维护大小为 \(k\)小根堆,堆顶是第 \(k\)
heapq 没有 key 元组表达多关键字
元组里放对象 中间插一个唯一整数当打破平局键,否则 TypeError
自定义类 只需实现 __lt__,但比元组慢 5 倍
删除任意元素 懒删除34 章
合并两个堆 拼接后重新 heapify\(O(n)\)
手写堆排序 只有教学意义,实战用 sorted()
\(n = 10^6\) heappush/heappop 绑成局部名
套路 识别特征
反悔贪心 「先全收下,超限了退掉最差的」
反悔贪心(收益型) 卖出后把卖价再推回堆,让差价可拼接
对顶堆 动态中位数
懒删除堆 元素会失效、需要跳过
堆优化 Dijkstra 带权最短路
heapq.merge 归并多个有序序列