Python heapq 模块详解

heapq 是 Python 标准库中用于实现堆队列（优先队列）的模块。它基于最小堆（min-heap）结构，常用于需要高效获取最小（或最大）元素的场景，如任务调度、Top-K 问题等。

什么是堆？

堆是一种特殊的完全二叉树，满足以下性质：

• 最小堆：父节点的值 ≤ 子节点的值，根节点为最小值。
• 最大堆：父节点的值 ≥ 子节点的值，根节点为最大值。

Python 的 heapq 模块默认实现的是最小堆。

常用函数

• heapq.heappush(heap, item)：将 item 推入堆中。
• heapq.heappop(heap)：弹出并返回堆中最小的元素。
• heapq.heapify(x)：将列表 x 原地转换为堆。
• heapq.heappushpop(heap, item)：先 push 再 pop，效率更高。
• heapq.heapreplace(heap, item)：先 pop 再 push。
• 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]
print(heapq.heappop(heap))  # 1
print(heap)           # [3, 5]

使用列表初始化堆

import heapq

data = [4, 7, 1, 9, 3]
heapq.heapify(data)
print(data)  # [1, 3, 4, 9, 7] —— 堆结构，非排序列表

实现最大堆技巧

由于 heapq 只支持最小堆，可通过插入负值模拟最大堆：

import heapq

max_heap = []
heapq.heappush(max_heap, -10)
heapq.heappush(max_heap, -5)
heapq.heappush(max_heap, -20)

# 弹出最大值（实际是最小的负数）
max_val = -heapq.heappop(max_heap)
print(max_val)  # 20

实用场景：Top-K 问题

import heapq

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

性能特点

• heappush 和 heappop：时间复杂度 O(log n)
• heapify：时间复杂度 O(n)
• nlargest/nsmallest：当 k 较小时效率高，k 接近 n 时不如排序