Python - heapq 模組介紹

Posted on

在 Python 中，有一個內建的 heapq 模組提供了最小堆積的函數。
在本文中，我們將會介紹 heapq 模組的用法。

heapq.heappush() 與 heapq.heappop()

heapq.heappush() 函數可以用來將元素加入到最小堆積中，而 heapq.heappop() 函數則是將最小堆積中的最小元素取出。
其中 heapq.heapify(x) 可以在線性時間 O(n) 時間內將 data 轉為 heap，且過程不會申請額外記憶體。image 106

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
import heapq

data = [63, 96, 39, 84, 12, 51, 20, 2, 34, 76]

heapq.heapify(data)
print(data)  # [2, 12, 20, 34, 63, 51, 39, 84, 96, 76]

heapq.heappush(data, 7)

print(data)  # [2, 7, 20, 34, 12, 51, 39, 84, 96, 76, 63]

print(heapq.heappop(data))  # 1
print(heapq.heappop(data))  # 2
print(heapq.heappop(data))  # 3

print(data)  # [20, 34, 39, 76, 63, 51, 96, 84]

heapq.heappushpop() 與 heapq.heapreplace()

heapq.heappushpop() 函數可以用來將元素加入到最小堆積中，並且取出最小元素。因此這個組合函式比呼叫 heappush() 之後呼叫 heappop() 更有效率。

1
2
3
4
5
6
7
8
9
10
11
import heapq

data = [63, 96, 39, 84, 12, 51, 20, 2, 34, 76]

heapq.heapify(data)

print(heapq.heappushpop(data, 7))  # 2
print(heapq.heappushpop(data, 2))  # 2
print(heapq.heappushpop(data, 3))  # 3

print(data)  # [7, 12, 20, 34, 63, 51, 39, 84, 96, 76]

heapq.heapreplace() 函數則是將最小堆積中的最小元素取出，並且加入新的元素。因此這個組合函式比呼叫 heappop() 之後呼叫 heappush() 更有效率。

1
2
3
4
5
6
7
8
9
10
11
import heapq

data = [63, 96, 39, 84, 12, 51, 20, 2, 34, 76]

heapq.heapify(data)

print(heapq.heapreplace(data, 7))  # 2
print(heapq.heapreplace(data, 2))  # 7
print(heapq.heapreplace(data, 3))  # 2

print(data)  # [3, 12, 20, 34, 63, 51, 39, 84, 96, 76]

heapq.nlargest() 與 heapq.nsmallest()

heapq.nlargest() 函數可以用來取出最大的 n 個元素，而 heapq.nsmallest() 函數則是取出最小的 n 個元素。

1
2
3
4
5
6
7
import heapq

data = [3, 12, 20, 34, 63, 51, 39, 84, 96, 76]
heapq.heapify(data)

print(heapq.nlargest(3, data))  # [97, 85, 73]
print(heapq.nsmallest(3, data))  # [3, 12, 20]

heapq.merge()

heapq.merge() 函數可以用來將多個已排序的 iterable 合併成一個已排序的 iterable。

1
2
3
4
5
6
7
8
9
10
import heapq

data1 = [1, 2, 3, 4, 5, 16, 32, 73, 85, 97]
data2 = [2, 3, 4, 5, 6, 17, 33, 74, 86, 98]
data3 = [3, 4, 5, 6, 7, 18, 34, 75, 87, 99]

for i in heapq.merge(data1, data2, data3):
    print(i, end=' ')

# 1 2 2 3 3 3 4 4 4 5 5 5 6 6 7 16 17 18 32 33 34 73 74 75 85 86 87 97 98 99

結論

heapq 模組提供了最小堆積的函數，可以讓我們在 Python 中更方便地使用最小堆積。

你可以在官方文件中找到更多 heapq 模組的函數，heapq — Heap queue algorithm。

也可以直接閱讀他的原始碼，heapq.py。

☺️

也許你也會想看看