堆 / 优先队列

一句话:堆是用数组下标隐式表达的完全二叉树。看着是树,内存是平的。每次 O(log n) 拿到当前最值,建堆只要 O(n)。

直觉 — 数组隐式表达的完全二叉树

堆的本质不是”一棵特殊的树”,而是一个数组外加一条下标规则:

  • 节点 i 的父亲在 (i-1)//2
  • 节点 i 的左右孩子在 2i+12i+2

没有指针、没有节点对象,树结构完全靠下标算出来。代价是堆必须是完全二叉树(从左到右紧密填满),这样数组才不会出现空洞。

这条规则带来三个性质,刚好对应所有堆操作:

操作 怎么做 复杂度
push 塞到数组末尾,沿父亲方向上浮(sift-up) O(log n)
pop 取出 heap[0],把末尾搬到 0,沿孩子方向下沉(sift-down) O(log n)
heapify 从最后一个非叶节点倒着 sift-down,Floyd 建堆 O(n) 而不是 O(n log n)

小根堆的不变量:每个节点 ≤ 它的两个孩子。注意只是父子有序,兄弟之间无序,所以堆不能用来排序后取第 k 个,只能反复 pop。

下面动画用 Top-K 场景把 push / sift-up / pop / sift-down 同时演示出来,左边是树形,右边是同一份数据的数组视角。重点盯着两个视图始终对应——这就是”数组模拟树”的全部内容。

模板 — Python heapq 基本操作

Python 标准库 heapq 只有小根堆。要大根堆?存 -x。这是面试题里最常见的坑。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
import heapq

heap = []
heapq.heappush(heap, x) # 入堆 O(log n)
top = heap[0] # 看堆顶最小值 O(1)
mn = heapq.heappop(heap) # 弹出最小值 O(log n)

# 大根堆:负数 trick
heapq.heappush(heap, -x)
mx = -heapq.heappop(heap)

# Floyd 建堆:原地 O(n),比逐个 push 的 O(n log n) 快
arr = [3, 1, 5, 2, 4, 6]
heapq.heapify(arr)

# 边 pop 边 push,比先 pop 后 push 快一倍
heapq.heapreplace(heap, x) # 先 pop 再 push
heapq.heappushpop(heap, x) # 先 push 再 pop

堆元素如果是元组,按字典序比较。遇到不可比的对象(如链表节点),加一个递增序号打破 tie:(val, idx, node)

4 类典型用法

算法题里的堆几乎只解决这四件事。认出场景比写代码重要。

1. Top-K — 大根堆找小,小根堆找大

要”前 K 大”用小根堆,堆大小恒为 K,堆顶就是”门槛”。新元素大于门槛才入堆。反过来要”前 K 小”就用大根堆。

1
2
3
4
5
6
7
def top_k_largest(nums, k):
heap = []
for x in nums:
heapq.heappush(heap, x)
if len(heap) > k:
heapq.heappop(heap) # 扔掉最小的,留下最大的 k 个
return heap # heap[0] = 第 k 大

时间 O(n log k),空间 O(k)。比排序 O(n log n) 更适合数据流。

2. 流式中位数 — 两个堆对顶

lo 是大根堆(放较小的一半),hi 是小根堆(放较大的一半)。维护 |len(lo) - len(hi)| ≤ 1,中位数永远在两个堆顶。

1
2
3
4
5
6
7
8
def add_num(num, lo, hi):
heapq.heappush(lo, -num) # 先放 lo
heapq.heappush(hi, -heapq.heappop(lo)) # lo 最大值倒给 hi
if len(hi) > len(lo): # 平衡
heapq.heappush(lo, -heapq.heappop(hi))

def median(lo, hi):
return -lo[0] if len(lo) > len(hi) else (-lo[0] + hi[0]) / 2

关键招数:先无脑塞 lo,再倒一个到 hi——保证 lo ≤ hi 这个不变量。

3. 合并 K 路有序 — 堆里放每路的当前指针

K 条链表/数组各拿出一个最小元素塞进堆,弹出最小那个,把它后面的元素补进堆。每次 O(log K)。

1
2
3
4
5
6
7
8
9
10
def merge_k_sorted(lists):
heap = [(lst[0], i, 0) for i, lst in enumerate(lists) if lst]
heapq.heapify(heap)
res = []
while heap:
val, i, j = heapq.heappop(heap)
res.append(val)
if j + 1 < len(lists[i]):
heapq.heappush(heap, (lists[i][j+1], i, j+1))
return res

总时间 O(N log K),N 是元素总数。

4. Dijkstra — 堆优化的最短路

朴素 Dijkstra 是 O(V²),用小根堆按距离取下一个节点,降到 O((V+E) log V)。堆里存 (dist, node),弹出时若 dist > dist[node] 就跳过(懒删除)。

1
2
3
4
5
6
7
8
9
10
11
12
13
def dijkstra(graph, src):
dist = {src: 0}
heap = [(0, src)]
while heap:
d, u = heapq.heappop(heap)
if d > dist[u]: # 过时的旧记录,丢掉
continue
for v, w in graph[u]:
nd = d + w
if nd < dist.get(v, float('inf')):
dist[v] = nd
heapq.heappush(heap, (nd, v))
return dist

懒删除的均摊复杂度仍是 O(log n):每个元素最多被入堆一次、被跳过一次,各 O(log n)。

经典题

  • 215. 数组中的第 K 个最大元素 — Top-K 模板,小根堆容量 K。
  • 295. 数据流的中位数 — 双堆对顶。
  • 23. 合并 K 个升序链表 — 堆里放节点,用序号打 tie。
  • 347. 前 K 个高频元素 — 先用 Counter 再上小根堆。
  • 502. IPO — 双堆:小根堆按 capital,大根堆按 profit。
  • 857. 雇佣 K 名工人的最低成本 — 排序 + 大根堆维护当前 K 人的工资和。
  • 703. 数据流中的第 K 大元素 — Top-K 在线版。
  • 1046. 最后一块石头的重量 — 大根堆模拟。

记忆法

把堆当成两层来记,就不会忘:

  1. 物理层:堆 = 数组 + 下标规则 parent=(i-1)//2, children=2i+1, 2i+2。所有”树”的操作最终落到数组下标的算术。
  2. 不变量层:小根堆只保证 parent ≤ children,兄弟无序。所以堆只能反复给你最值,不能给你”第二大”。
  3. 场景层:看到下面这些关键词条件反射上堆:
    • “前 K 个 / 第 K 大 / 第 K 小” → Top-K(一个堆)
    • “数据流中位数 / 滑动窗口中位数” → 双堆对顶
    • “K 路 / K 个有序” → K 路归并
    • “最短路 / 最小代价扩展” → Dijkstra / 类 BFS 加权版

Python 侧只记三行就够:heappush / heappop / heap[0],大根堆负号一翻,heapify 一次 O(n)。

思考题

为什么 Top K 用最小堆而不是最大堆?

最小堆的堆顶是最小值,可以快速判断新元素是否比当前第 K 大更大。如果用最大堆,堆顶是最大值,没法高效淘汰最小的元素。最小堆保证堆中始终是最大的 K 个,堆顶恰好是”门槛”。

双堆中为什么"先放 lo 再倒到 hi"能保证 lo ≤ hi?

先放 lo(最大堆),再把 lo 的最大值倒到 hi(最小堆)。这保证了从 lo 转移到 hi 的元素一定 ≥ lo 中剩余的所有元素。因此 lo 的最大值 ≤ hi 的最小值,即 lo 所有元素 ≤ hi 所有元素。

Floyd 建堆为什么是 O(n) 而不是 O(n log n)?

逐个 push 是 O(n log n):每次上浮最坏 log n。Floyd 建堆从最后一个非叶节点倒着 sift-down,叶子那一层(n/2 个节点)完全不动,倒数第二层每个最多下沉 1 步,再上一层最多 2 步……求和 Σ (n/2^(h+1)) · h 收敛到 O(n)。直觉:大部分节点都在底层,它们几乎不用动

Dijkstra 里的"懒删除"为什么不会爆复杂度?

每条边最多让目标节点入堆一次,所以堆里最多 O(E) 个条目。每个条目要么被有效处理,要么被 d > dist[u] 跳过——两种情况各是 O(log E)。总复杂度 O((V+E) log V),和理论上界一致。

LeetCode 练习

热身:215、703、347、1046。

进阶:23、295、378、502、857。