Python heapq 模块详解

什么是 heapq？

heapq 是 Python 标准库中的一个模块，用于实现堆（heap）数据结构。 它提供了一系列函数，可以在列表上原地维护一个最小堆（即堆顶元素为最小值）。 虽然名为“heap”，但 heapq 并不提供堆类，而是通过函数操作普通列表。

常用函数一览

• heapq.heapify(x)：将列表 x 转换为堆，时间复杂度 O(n)。
• heapq.heappush(heap, item)：向堆中添加元素。
• heapq.heappop(heap)：弹出并返回堆中最小元素。
• heapq.heappushpop(heap, item)：先 push 再 pop，效率更高。
• heapq.heapreplace(heap, item)：先 pop 再 push。
• heapq.merge(*iterables)：合并多个已排序的可迭代对象。
• heapq.nlargest(n, iterable)：返回最大的 n 个元素。
• heapq.nsmallest(n, iterable)：返回最小的 n 个元素。

基础使用示例

import heapq

# 创建一个空堆
heap = []

# 添加元素
heapq.heappush(heap, 5)
heapq.heappush(heap, 1)
heapq.heappush(heap, 3)

print(heap)  # [1, 5, 3] —— 注意：不是完全排序，但满足堆性质

# 弹出最小元素
min_val = heapq.heappop(heap)
print(min_val)  # 1
print(heap)     # [3, 5]

从列表构建堆

import heapq

data = [10, 3, 7, 1, 9]
heapq.heapify(data)
print(data)  # [1, 3, 7, 10, 9]

注意：heapify() 是原地操作，且只保证堆性质，不保证完全有序。

实际应用场景

1. 实现优先队列

由于堆天然支持快速获取最小（或最大）元素，常用于任务调度、Dijkstra 算法等场景。

2. 找出 Top-K 元素

import heapq

scores = [85, 92, 78, 96, 88, 91]
top3 = heapq.nlargest(3, scores)
print(top3)  # [96, 92, 91]

3. 合并多个有序流

import heapq

list1 = [1, 4, 7]
list2 = [2, 5, 8]
list3 = [3, 6, 9]

merged = list(heapq.merge(list1, list2, list3))
print(merged)  # [1, 2, 3, 4, 5, 6, 7, 8, 9]

注意事项

• heapq 默认实现的是最小堆。若需最大堆，可对元素取负值。
• 堆中的元素必须支持比较操作（如数字、字符串等）。
• 不要手动修改堆列表后不调用 heapify() 或相关函数，否则会破坏堆结构。