heapq 模块实现了适用于python列表的最小堆排序算法，
该模块提供 heappush 和 heapify 两种方法来构建最小堆，heappop 方法可以删除堆顶元素，
nlargest 和 nsmallest 可以分别查看最大的几个元素和最小的几个元素，关于堆，
如果不了解其概念，可以先看下面这段介绍。
堆的概念
假定在数据记录中存在一个能够标识数据记录的数据项，并可依据该数据项对数据进行组织，
则称此数据项为关键码（ key ）。
如果有一个关键码集合 K = {k0 , k1 , k2 , k3 , ... kn-1 } ，
把它的所有元素按照完全二叉树的顺序存储方式存放在一个一维数组中，
并且满足 ki ≤ k2i+1 且 ki ≤ k2i+2 （或者 ki ≥ k2i+1 且 ki ≥ k2i+2 ） ,
i = 0, 1,2, ... (n-2)/2 ，则称这个集合为最小堆(最大堆)。
在最小堆中，父节点的关键码小于等于它的左右子女的关键码，最大堆中， 父节点的关键码大于等于左右子女的关键码。
构建最小堆
使用 heappush 方法创建
In [2]:
import heapq
In [3]:
random_lst = [5, 7, 2, 1, 6, 10, 8, 9]
heap_lst = []
In [4]:
for item in random_lst:
heapq.heappush(heap_lst, item)
In [5]:
print(heap_lst)
[1, 2, 5, 7, 6, 10, 8, 9]
使用 heappush 方法，将混乱的数据依次放入到 heap_lst 中，就会得到一个最小堆，程序输出结果为：
[1, 2, 5, 7, 6, 10, 8, 9]如果你不熟悉堆的概念，恐怕难以从结果上看出门道， 不妨编写一个输出堆的函数，以更加立体的方式展示堆：
In [6]:
import math
from io import StringIO
In [7]:
def print_heap(heap, width=36, fill=' '):
output = StringIO()
last_row = -1
for index, item in enumerate(heap):
# 计算元素在哪一行
if index:
row = int(math.floor(math.log(index + 1, 2)))
else:
row = 0
# 换行操作
if row != last_row:
output.write('\n')
columns = 2 ** row # 计算一行有多少个数据
col_width = int(math.floor(width / columns))
output.write(str(item).center(col_width, fill))
last_row = row
print(output.getvalue())
In [8]:
print_heap(heap_lst)
1
2 5
7 6 10 8
9
程序输出结果为：
1
2 5
7 6 10 8
9 In [9]:
random_lst = [5, 7, 2, 1, 6, 10, 8, 9]
heapq.heapify(random_lst)
print_heap(random_lst)
1
5 2
7 6 10 8
9
In [10]:
random_lst = [9, 7, 5, 2, 1, 6, 10, 8]
heapq.heapify(random_lst)
print_heap(random_lst)
print("-"*20)
1
2 5
8 7 6 10
9
--------------------
In [11]:
heapq.heappop(random_lst)
print_heap(random_lst)
2
7 5
8 9 6 10
程序输出结果。
1
2 5
8 7 6 10
9
--------------------
2
7 5
8 9 6 10 In [12]:
random_lst = [9, 7, 5, 2, 1, 6, 10, 8]
heapq.heapify(random_lst)
print_heap(random_lst)
print("-"*20)
1
2 5
8 7 6 10
9
--------------------
In [13]:
heapq.heapreplace(random_lst, 4)
print_heap(random_lst)
2
4 5
8 7 6 10
9
程序输出结果：
1
2 5
8 7 6 10
9
--------------------
2
4 5
8 7 6 10
9 In [14]:
random_lst = [9, 7, 5, 2, 1, 6, 10, 8, 4, 3]
heap_lst = []
In [15]:
for item in random_lst[:3]:
heapq.heappush(heap_lst, item)
In [16]:
for item in random_lst[3:]:
if item > heap_lst[0]:
heapq.heapreplace(heap_lst, item)
In [17]:
print(heap_lst)
[8, 9, 10]
程序输出结果。
[8, 9, 10]对于一个给定的最小堆，nlargest和 nsmallest 可以分别查看最大的几个元素和最小的几个元素。
In [18]:
random_lst = [9, 7, 5, 2, 1, 6, 10, 8, 4, 3]
heapq.heapify(random_lst)
In [19]:
print("最大的3个元素", heapq.nlargest(3, random_lst))
最大的3个元素 [10, 9, 8]
In [20]:
print("最小的3个元素", heapq.nsmallest(3, random_lst))
最小的3个元素 [1, 2, 3]
程序输出结果
最大的3个元素 [10, 9, 8]
最小的3个元素 [1, 2, 3]In [21]:
import heapq
import random
In [22]:
lst = []
In [23]:
for i in range(5):
l = random.sample(range(1, 100), 5)
l.sort()
print(l)
lst.append(l)
[13, 62, 66, 70, 92] [3, 5, 34, 52, 74] [23, 35, 42, 87, 94] [25, 33, 50, 84, 94] [17, 63, 64, 83, 86]
In [26]:
print("合并后")
合并后
In [27]:
merge = [i for i in heapq.merge(*lst)]
print(merge)
[2, 3, 5, 6, 8, 13, 14, 17, 19, 19, 23, 25, 25, 27, 31, 33, 34, 35, 36, 37, 37, 38, 41, 42, 44, 45, 48, 50, 52, 62, 62, 63, 64, 65, 66, 68, 70, 71, 74, 78, 78, 79, 83, 84, 86, 87, 92, 94, 94, 94]
程序输出结果。
[3, 10, 13, 20, 21, 43, 48, 53, 54, 56, 60, 62, 63, 63, 69, 70, 75, 77, 80, 81, 83, 89, 93, 95, 97]