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庫的用法介紹