从算法到人工智能 · 第 13 课：堆与优先队列——永远知道"最大/最小"

想象一个场景：医院急诊室。病人源源不断地来，但病情重的必须先看，不能简单按先来后到排队。你需要一个数据结构，它能做到——随时把"最急的那个"拎出来，新病人来了随时插进去。

这个"永远知道谁最大/最小"的数据结构，就是堆（heap）。它是优先队列（priority queue）的实现方式，也是很多算法（topK、Dijkstra、任务调度）背后的功臣。

一、优先队列是什么：会"插队"的队列

第 2 课我们学过普通队列：先进先出（FIFO）。但现实中经常需要"按优先级出队"，于是有了优先队列：

每次出队，弹出的不是"最先来的"，而是优先级最高（最大或最小）的。

import heapq   # Python 自带的小顶堆

pq = []
heapq.heappush(pq, 5)   # 插入
heapq.heappush(pq, 1)
heapq.heappush(pq, 3)
print(heapq.heappop(pq))   # 1 —— 最小先出
print(heapq.heappop(pq))   # 3
print(heapq.heappop(pq))   # 5

这就是优先队列：插入随便插，弹出永远弹最小（或最大）的。

二、堆是什么：一棵"矮胖"的完全二叉树

堆在概念上是一棵完全二叉树，但有个严格的约束——堆序性：

• 小顶堆：每个节点的值 ≤ 它的所有孩子（根是最小值）
• 大顶堆：每个节点的值 ≥ 它的所有孩子（根是最大值）

小顶堆：
        1          ← 根永远最小
       / \
      3   5
     / \
    9   7

堆有两个"神奇"性质，是它高效的根本：

性质 1：根永远是极值

所以取最小/最大，O(1) 直接看根——不用扫描。

性质 2：用数组就能存，不用指针

还记得第 12 课的"下标 2i+1 / 2i+2"吗？堆就是用这个规则存在数组里：

数组: [1, 3, 5, 9, 7]
下标:  0  1  2  3  4

节点 i 的左孩子 = 2i+1，右孩子 = 2i+2，父节点 = (i-1)//2

所以 heapq 底层就是一个普通 list，用下标换算父子关系，不用建 TreeNode 类。

三、堆的两个核心操作

插入（push）——从底部"上浮"

新元素先放到数组末尾（完全二叉树的最右下），然后和父节点比：如果它更小（小顶堆），就往上"浮"，直到满足堆序。

插入 2：
        1
       / \
      3   5
     / \ /
    9  7 2   ← 2 比父节点 5 小，上浮
          ↓
        1
       / \
      3   2
     / \ /
    9  7 5   ← 2 再和 1 比，不小了，停

弹出（pop）——根和末尾交换再"下沉"

弹出根（最小值）后，把数组最后一个元素挪到根的位置，然后和较小的孩子比，往下"沉"，直到满足堆序。

弹出 1：把 7 挪到根
        7
       / \
      3   5
     /
    9       ← 7 比孩子 3 大，下沉
          ↓
        3
       / \
      7   5
     /
    9       ← 7 比 9 小，停

上浮和下沉，都是沿着树走 O(log n) 层。 所以：

四、Python 里用 heapq：三个必会操作

import heapq

nums = [5, 1, 3, 9, 7]

# 1. 原地建堆（注意：是小顶堆）
heapq.heapify(nums)
print(nums)   # [1, 3, 5, 9, 7]  根是最小

# 2. 弹出最小
print(heapq.heappop(nums))   # 1

# 3. 插入
heapq.heappush(nums, 0)
print(nums[0])   # 0，新的最小

想要"大顶堆"怎么办？

Python 的 heapq 只提供小顶堆。想要大顶堆，把值取负存进去，弹出时再取负回来：

import heapq
max_heap = []
for x in [5, 1, 3, 9, 7]:
    heapq.heappush(max_heap, -x)   # 存负数
print(-heapq.heappop(max_heap))    # 9 —— 最大

记忆口诀：Python 堆 = 小顶堆，要最大就存负数。

五、经典应用：Top K 问题

面试最高频的堆应用题：给你一堆数，找最大的 K 个。

朴素想法：全排序，取前 K —— O(n log n)。
堆的做法：维护一个大小为 K 的小顶堆，遍历所有数——如果当前数比堆顶（K 个里最小的）大，就把它换进去。最后堆里就是最大的 K 个。

import heapq

def top_k(nums, k):
    heap = []
    for x in nums:
        if len(heap) < k:
            heapq.heappush(heap, x)      # 堆还没满，直接进
        elif x > heap[0]:
            heapq.heapreplace(heap, x)   # 弹出最小，塞入 x
    return heap

nums = [3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5]
print(sorted(top_k(nums, 3)))   # [6, 9, 5] → 最大的三个

# 更简洁：用 heapq.nlargest 一行搞定
print(heapq.nlargest(3, nums))   # [9, 6, 5]

为什么用小顶堆而不是大顶堆？ 关键：小顶堆的堆顶是"当前 K 个里最小的"，新数只要和它比，就能 O(1) 判断"该不该进来"。用大顶堆反而要一直看最大，没法快速淘汰。

复杂度：O(n log k)。当 k 远小于 n 时，比全排序 O(n log n) 快得多。

六、堆排序：用堆给数组排序

第 8 课提过堆排序，现在补全。思路：建堆 → 反复弹根。

import heapq

def heap_sort(nums):
    heapq.heapify(nums)          # O(n) 建堆
    result = []
    while nums:
        result.append(heapq.heappop(nums))   # 每次弹最小
    return result

nums = [5, 1, 3, 9, 7, 2]
print(heap_sort(nums))   # [1, 2, 3, 5, 7, 9]

堆排序复杂度：O(n log n)，而且原地、不稳定。它和快排、归并都是 O(n log n)，但常数更大，所以实际排序常用快排；堆的舞台在"动态取极值"，而不是纯排序。

七、复杂度小结

八、动手时间 🎯

实验 1：亲眼看看堆的"根永远是极值"

import heapq

heap = []
import random
for _ in range(5):
    heapq.heappush(heap, random.randint(1, 100))

print("整个堆:", heap)
print("根（最小）:", heap[0])   # 永远是全堆最小

你会看到：不管 push 什么顺序，heap[0] 永远是最小的那个。

实验 2：Top K 性能对比（堆 vs 全排序）

import heapq, time, random

nums = [random.randint(1, 100000) for _ in range(1000000)]
k = 10

# 方法1：全排序
start = time.time()
sorted(nums)[-k:]
t1 = time.time() - start

# 方法2：堆
start = time.time()
heapq.nlargest(k, nums)
t2 = time.time() - start

print(f"全排序: {round(t1,4)}s | 堆: {round(t2,4)}s")

体会：100 万个数据找 Top 10，堆明显更快——因为它只维护 10 个元素，不用全排。

实验 3：用堆实现"任务调度"（带优先级）

import heapq

# 任务：(优先级, 任务名)，数字越小越优先
tasks = [(3, "写周报"), (1, "处理服务器告警"), (2, "回复邮件")]
heapq.heapify(tasks)

while tasks:
    priority, name = heapq.heappop(tasks)
    print(f"处理: {name}（优先级 {priority}）")
# 先"服务器告警"，再"邮件"，最后"周报"

实验 4（挑战）：合并 K 个有序链表

经典题：有 K 个已经排好序的列表，合并成一个有序列表。用堆维护"每个列表当前最小的那个"：

import heapq

def merge_k_lists(lists):
    heap = []
    # 每个列表的第一个元素进堆：(值, 列表下标, 元素在列表内的下标)
    for i, lst in enumerate(lists):
        if lst:
            heapq.heappush(heap, (lst[0], i, 0))

    result = []
    while heap:
        val, i, j = heapq.heappop(heap)
        result.append(val)
        if j + 1 < len(lists[i]):      # 该列表还有下一个，进堆
            heapq.heappush(heap, (lists[i][j + 1], i, j + 1))
    return result

lists = [[1, 4, 5], [1, 3, 4], [2, 6]]
print(merge_k_lists(lists))   # [1, 1, 2, 3, 4, 4, 5, 6]

体会：堆在这里扮演"裁判"——始终盯着每个列表最小的那个，谁小谁出列。这就是"多路归并"的核心思想，也是第 16 课 Dijkstra 最短路径算法的底层机制。

九、小结

1. 堆 = 永远知道最大/最小的完全二叉树，根是极值（O(1)），插入弹出都是 O(log n)，用数组存（2i+1/2i+2）。
2. Python 的 heapq 是小顶堆，要最大就存负数；核心操作 heapify/heappush/heappop。
3. 堆的最佳舞台是"动态取极值"——Top K（小顶堆维护 K 个）、合并 K 个有序序列、任务调度、Dijkstra，都是它的主场。

先别急着往后翻，把 Top K 的"堆 vs 全排序"跑一遍，亲手感受 O(n log k) 对 O(n log n) 的碾压，堆的价值就刻在你脑子里了。