Python中的heapq模块

Python中的heapq模块

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库的用法介绍