Python中的heapq模块
文章目录
- Python中的heapq模块
- 1.heapq的方法
- 2.使用heapq创建堆
- 3.使用heapq实现堆排序
- 4.获取堆中的前n个最大值或最小值
- Reference
heapq模块实现了堆队列的算法，即优先队列算法。heapq其实是实现了一种小顶堆，所以使用pop()方法返回的是当前堆中的最小元素。
1.heapq的方法
| 方法 | 功能 |
|---|---|
heapq.heappush(heap, item) | 将item的值加入到heap中，保持堆的不变性 |
heapq.heappop(heap) | 弹出并返回heap中的最小值，保持堆的不变性。 |
heapq.heappushpop(heap, item) | 将item放入堆中，然后弹出并返回heap中的最小元素，这个操作比先调用heappush()再调用heappop()效率更高。 |
heapq.heapify(x) | 将list x转换成堆 |
heapq.heapreplace(heap, item) | 弹出并返回 heap 中最小的一项，同时推入新的 item。 |
heapq.nlargest(n, iterable, key=None) | 从 iterable 所定义的数据集中返回前 n 个最大元素组成的列表。 |
heapq.nlargest(n, iterable, key=None) | 从 iterable 所定义的数据集中返回前 n 个最小元素组成的列表。 |
heapq.merge(*iterables, key=None, reverse=False) | 将多个已排序的输入合并为一个已排序的输出 |
2.使用heapq创建堆
有两种方法可以用于创建堆，第一种是直接使用方法heapq.heapify(iterable)，直接将可迭代的对象转换成小顶堆。第二种方法是使用heapq.push(heap, item)将元素手动放入指定的heap中。
import heapq
array = [5, 7, 9, 0, 3, 2, 1, 6, 4, 8]
# 1.使用heapq.push来创建
heap = []
for num in array:heapq.heappush(heap, num)
print('array:', array)
print('heap:', heap)
# 2.使用heapify来创建
heapq.heapify(array)
print('array:', array)
array: [5, 7, 9, 0, 3, 2, 1, 6, 4, 8]
heap: [0, 3, 1, 4, 5, 9, 2, 7, 6, 8]
array: [0, 3, 1, 4, 5, 2, 9, 6, 7, 8]
特别注意的是，堆元素可以为元组，这有利于以下做法——在被跟踪的主记录旁边添一个额外的值（例如任务的优先级）用于互相比较，我们只需要将排序的值放在元组的第一个位置即可：
import heapq
heap = []
heapq.heappush(heap, (5, 'Alex'))
heapq.heappush(heap, (2, 'Ben'))
heapq.heappush(heap, (0, 'David'))
heapq.heappush(heap, (1, 'Elon'))
print('heap:', heap)
heapq.heappop(heap)
print('heap:', heap)
heap: [(0, 'David'), (1, 'Elon'), (2, 'Ben'), (5, 'Alex')]
heap: [(1, 'Elon'), (5, 'Alex'), (2, 'Ben')]
这里我们按照tuple中第一元素，即这个数字来进行比较，构成堆，我们弹出的最小的元素是值为0的David。
import heapq
heap = []
heapq.heappush(heap, ('Alex', 5))
heapq.heappush(heap, ('Ben', 2))
heapq.heappush(heap, ('David', 0))
heapq.heappush(heap, ('Elon', 1))
print('heap:', heap)
heapq.heappop(heap)
print('heap:', heap)
heap: [('Alex', 5), ('Ben', 2), ('David', 0), ('Elon', 1)]
heap: [('Ben', 2), ('Elon', 1), ('David', 0)]
如果我们反过来使用名字来排序，构成堆，我们弹出的最小元素是ASCII码最小的A即Alex。
3.使用heapq实现堆排序
我们可以将待排序的数据构建成一个小顶堆，每次从堆顶弹出数据，收集弹出的数据，这样我们就可以获得一个排完序的序列。
import heapq
array = [5, 7, 9, 0, 3, 2, 1, 6, 4, 8]
heapq.heapify(array)
res = [heapq.heappop(array) for _ in range(len(array))]
print('heap sort result:', res)
heap sort result: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
4.获取堆中的前n个最大值或最小值
import heapq
array = [5, 7, 9, 0, 3, 2, 1, 6, 4, 8]
print('3 largest value:',heapq.nlargest(3, array))
print('3 smallest value:',heapq.nsmallest(3, array))
3 largest value: [9, 8, 7]
3 smallest value: [0, 1, 2]
Reference
heapq官方文档
Python heapq库的用法介绍
)
