Python中的heapq模块_python heapq

Python中的heapq模块

heapq模块实现了堆队列的算法，即优先队列算法。heapq其实是实现了一种小顶堆，所以使用pop()方法返回的是当前堆中的最小元素。

1.heapq的方法

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)
1
2
3
4
5
6
7
8
9
10
11
'
运行
运行

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]
1
2
3

特别注意的是，堆元素可以为元组，这有利于以下做法——在被跟踪的主记录旁边添一个额外的值（例如任务的优先级）用于互相比较，我们只需要将排序的值放在元组的第一个位置即可：

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)
1
2
3
4
5
6
7
8
9
'
运行
运行

heap: [(0, 'David'), (1, 'Elon'), (2, 'Ben'), (5, 'Alex')]
heap: [(1, 'Elon'), (5, 'Alex'), (2, 'Ben')]
1
2

这里我们按照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)
1
2
3
4
5
6
7
8
9
'
运行
运行

heap: [('Alex', 5), ('Ben', 2), ('David', 0), ('Elon', 1)]
heap: [('Ben', 2), ('Elon', 1), ('David', 0)]
1
2

如果我们反过来使用名字来排序，构成堆，我们弹出的最小元素是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)
1
2
3
4
5
'
运行
运行

heap sort result: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
1

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))
1
2
3
4
'
运行
运行

3 largest value: [9, 8, 7]
3 smallest value: [0, 1, 2]
1
2

Reference

heapq官方文档
Python heapq库的用法介绍