05_优先队列

优先队列

模板

import heapq

pq = []
heapq.heapify(pq)            # 把已有列表原地转成堆，O(n)

heapq.heappush(pq, x)        # 插入元素
smallest = pq[0]             # 获取最小元素（不弹出）
smallest = heapq.heappop(pq) # 弹出并返回最小值

方法	说明
heapq.heappush(pq, x)	插入元素 x
heapq.heappop(pq)	弹出并返回最小值
heapq.heapreplace(pq, x)	先弹最小值再插入 x（一次下沉，比 pop+push 快）
heapq.heappushpop(pq, x)	先插入 x 再弹最小值
heapq.heapify(pq)	将列表原地转换为堆
heapq.nlargest(k, a) / heapq.nsmallest(k, a)	最大的 k 个 / 最小的 k 个

大顶堆

Python 的 heapq 只有小顶堆，大顶堆用取负实现：

heapq.heappush(heap, -x)      # 存入负值
x = -heapq.heappop(heap)      # 取出时再取负，得到原值

堆元素是元组时

按元组第一关键字排序（第一关键字相同时看第二关键字），常用于「按 (距离, 节点) 取最近」：

heapq.heappush(pq, (dist, node))
d, u = heapq.heappop(pq)

解释

用途：$O(\log N)$ 插入 / 删除 / 获取最高优先级元素。凡是「反复取当前最大/最小」的场景（贪心、多路归并、Top K）都用它。

优先队列（Priority Queue / Heap）比有序集合更简单更快，能用堆就不要维护排序数组。

使用 heapq 的三个注意点：

1. Python（与 C++ 的 priority_queue 相反）删除和获取的是最小元素
2. heapq 不是封装好的类，而是直接操作传入的列表——传同一个列表就是在操作同一个堆
3. 堆内部列表不保证有序，print(pq) 看到的是数组布局而非排序结果，只有 pq[0] 保证是最小值

例题

例题 1：Room Allocation（小顶堆 + 贪心）

来源：CSES - Room Allocation

给定 n 个客户的到达和离开时间，求所需的最少房间数，并输出每人分配的房间编号。

思路：

1. 按到达时间排序所有客户
2. 维护一个小顶堆，存储（已入住客户的离开时间, 房间号）
3. 对于每个客户：
   • 若堆顶的离开时间 < 新客户的到达时间，说明有房间空出，heapreplace 复用该房间
   • 否则所有房间都满了，开新房间

import heapq

n = int(input())
timetable = []
for i in range(n):
    a, b = map(int, input().split())
    timetable.append((a, b, i))
timetable.sort(key=lambda x: x[0])

heap = []
ans = 0
results = [-1] * n

for a, b, i in timetable:
    if heap and a > heap[0][0]:
        r = heap[0][1]
        heapq.heapreplace(heap, (b, r))   # 复用房间：弹最早的离开时间，放新客户的
    else:
        ans += 1
        r = ans
        heapq.heappush(heap, (b, r))      # 新开房间
    results[i] = r

print(ans)
print(*results)

复杂度：时间 O(n log n)，空间 O(n)

适用场景

1. 贪心算法：每次选择最小/最大的元素
2. 多路归并：合并多个有序序列
3. 求 Top K：维护大小为 K 的堆（配 nlargest / nsmallest）
4. 任务调度：按优先级处理任务
5. Dijkstra 最短路：每次取出距离最小的节点（见《06Graphs/07图论》）