Python 中的 heapq 模块简介

heapq 模块简介

heapq 模块提供了实现堆队列算法的函数。堆是一种特殊的树形数据结构，每个父节点都小于或等于（最小堆）或大于或等于（最大堆）其孩子节点。heapq 实现了最小堆，这意味着堆顶元素总是最小的。

主要函数

1. heapq.heappush(heap, item)
   • 将一个元素 item 推入堆中，保持堆的性质。
2. heapq.heappop(heap)
   • 弹出堆顶元素（即最小的元素），并返回它。如果堆为空，则引发 IndexError。
3. heapq.heappushpop(heap, item)
   • 将 item 推入堆中，然后弹出并返回堆顶元素。这个操作比先调用 heappush() 再调用 heappop() 更有效率。
4. heapq.heapify(x)
   • 将列表 x 转换为堆，就地（in-place）改变列表 x。这意味着不会返回一个新的堆，而是修改输入的列表，使其成为堆。
5. heapq.heapreplace(heap, item)
   • 弹出堆顶元素，并将 item 推入堆中。这个操作比先调用 heappop() 再调用 heappush() 更有效率。
6. *heapq.merge(iterables)
   • 将多个已排序的输入迭代器合并为一个迭代器，产生排序的输出。这类似于 sorted(itertools.chain(*iterables))，但合并操作是懒惰的，即它不会一次性产生整个排序列表，而是返回一个迭代器，按需产生排序的元素。
7. heapq.nlargest(n, iterable, key=None)
   • 返回迭代器 iterable 中最大的 n 个元素组成的列表。
8. heapq.nsmallest(n, iterable, key=None)
   • 返回迭代器 iterable 中最小的 n 个元素组成的列表。

使用示例

示例 1：使用 heappush 和 heappop

import heapq  
  
heap = []  
heapq.heappush(heap, 3)  
heapq.heappush(heap, 1)  
heapq.heappush(heap, 4)  
  
print(heapq.heappop(heap))  # 输出: 1  
print(heapq.heappop(heap))  # 输出: 3
'
运行
运行

示例 2：使用 heapify

import heapq  
  
lst = [3, 1, 4]  
heapq.heapify(lst)  
  
print(lst)  # 输出: [1, 3, 4]（堆的结构，但顺序可能因实现而异）  
print(heapq.heappop(lst))  # 输出: 1
'
运行
运行

示例 3：使用 merge

import heapq  
import itertools  
  
a = [1, 3, 5, 7]  
b = [2, 4, 6, 8]  
  
for i in heapq.merge(a, b):  
    print(i)  # 输出: 1 2 3 4 5 6 7 8
'
运行
运行

注意事项

• heapq 模块实现的是最小堆，如果你需要最大堆，可以通过将值取负来实现，但这样会增加一些额外的计算开销。
• heapq 模块仅提供了堆的部分操作，并没有实现完整的堆数据结构（如二叉堆或斐波那契堆）。如果你需要更复杂的堆操作，可能需要使用其他数据结构或库。
• 堆是一个部分有序的树形结构，只有父节点和子节点之间满足大小关系，兄弟节点之间不一定有序。因此，你不能通过直接访问索引来获取排序的元素（如 heap[0]），因为堆的内部结构可能因操作而改变。你应该始终使用 heappop() 来获取堆顶元素。