Python3標准庫：heapq堆排序算法

1. heapq堆排序算法

堆(heap)是一個樹形數據結構，其中子節點與父節點有一種有序關系。二叉堆(binary heap)可以使用一個有組織的列表或數組表示，其中元素N的子元素位於2*N+1和2*N+2(索引從0開始)。這種布局允許原地重新組織堆，從而不必再添加或刪除元素時重新分配大量內存。

最大堆(max-heap)確保父節點大於或等於其兩個子節點。最小堆(min-heap)要求父節點小於或等於其子節點。Python的heapq模塊實現了一個最小堆。

1.1 創建堆

創建堆有兩種基本方式：heappush()和heapify()。

import heapq
import math
from io import StringIO

data = [19, 9, 4, 10, 11]

def show_tree(tree, total_width=36, fill=' '):
    """Pretty-print a tree."""
    output = StringIO()
    last_row = -1
    for i, n in enumerate(tree):
        if i:
            row = int(math.floor(math.log(i + 1, 2)))
        else:
            row = 0
        if row != last_row:
            output.write('\n')
        columns = 2 ** row
        col_width = int(math.floor(total_width / columns))
        output.write(str(n).center(col_width, fill))
        last_row = row
    print(output.getvalue())
    print('-' * total_width)
    print()

heap = []
print('random :', data)
print()

for n in data:
    print('add {:>3}:'.format(n))
    heapq.heappush(heap, n)
    show_tree(heap)

使用heappush()，從數據源增加新元素時會保持元素的堆排序順序。

如果數據已經在內存中，那么使用heapify()原地重新組織列表中的元素會更高效。

import heapq
import math
from io import StringIO

data = [19, 9, 4, 10, 11]

def show_tree(tree, total_width=36, fill=' '):
    """Pretty-print a tree."""
    output = StringIO()
    last_row = -1
    for i, n in enumerate(tree):
        if i:
            row = int(math.floor(math.log(i + 1, 2)))
        else:
            row = 0
        if row != last_row:
            output.write('\n')
        columns = 2 ** row
        col_width = int(math.floor(total_width / columns))
        output.write(str(n).center(col_width, fill))
        last_row = row
    print(output.getvalue())
    print('-' * total_width)
    print()

print('random    :', data)
heapq.heapify(data)
print('heapified :')
show_tree(data)

如果按堆順序一次一個元素地構建列表，那么結果與構建一個無序列表再調用heapify()是一樣的。

1.2 訪問堆內容

一旦堆已經被正確組織，則可以使用heappop()刪除有最小值的元素。

import heapq
import math
from io import StringIO

data = [19, 9, 4, 10, 11]

def show_tree(tree, total_width=36, fill=' '):
    """Pretty-print a tree."""
    output = StringIO()
    last_row = -1
    for i, n in enumerate(tree):
        if i:
            row = int(math.floor(math.log(i + 1, 2)))
        else:
            row = 0
        if row != last_row:
            output.write('\n')
        columns = 2 ** row
        col_width = int(math.floor(total_width / columns))
        output.write(str(n).center(col_width, fill))
        last_row = row
    print(output.getvalue())
    print('-' * total_width)
    print()

print('random    :', data)
heapq.heapify(data)
print('heapified :')
show_tree(data)
print()

for i in range(2):
    smallest = heapq.heappop(data)
    print('pop    {:>3}:'.format(smallest))
    show_tree(data)

這個例子是由標准庫文檔改寫的，其中使用heapify()和heappop()對一個數字隊列進行排序。

如果希望在一個操作中刪除現有元素並替換為新值，則可以使用heapreplace()。

import heapq
import math
from io import StringIO

data = [19, 9, 4, 10, 11]

def show_tree(tree, total_width=36, fill=' '):
    """Pretty-print a tree."""
    output = StringIO()
    last_row = -1
    for i, n in enumerate(tree):
        if i:
            row = int(math.floor(math.log(i + 1, 2)))
        else:
            row = 0
        if row != last_row:
            output.write('\n')
        columns = 2 ** row
        col_width = int(math.floor(total_width / columns))
        output.write(str(n).center(col_width, fill))
        last_row = row
    print(output.getvalue())
    print('-' * total_width)
    print()

heapq.heapify(data)
print('start:')
show_tree(data)

for n in [0, 13]:
    smallest = heapq.heapreplace(data, n)
    print('replace {:>2} with {:>2}:'.format(smallest, n))
    show_tree(data)

通過原地替換元素，就這樣可以維持一個固定大小的堆，如按優先級排序的作業隊列。

1.3 堆的數據極值

heapq還包括兩個檢查可迭代對象(iterable)的函數，可以查找其中包含的最大或最小值的范圍。

import heapq
import math
from io import StringIO

data = [19, 9, 4, 10, 11]

def show_tree(tree, total_width=36, fill=' '):
    """Pretty-print a tree."""
    output = StringIO()
    last_row = -1
    for i, n in enumerate(tree):
        if i:
            row = int(math.floor(math.log(i + 1, 2)))
        else:
            row = 0
        if row != last_row:
            output.write('\n')
        columns = 2 ** row
        col_width = int(math.floor(total_width / columns))
        output.write(str(n).center(col_width, fill))
        last_row = row
    print(output.getvalue())
    print('-' * total_width)
    print()

print('all       :', data)
print('3 largest :', heapq.nlargest(3, data))
print('from sort :', list(reversed(sorted(data)[-3:])))
print('3 smallest:', heapq.nsmallest(3, data))
print('from sort :', sorted(data)[:3])

只有當n值(n>1)相對小時使用nlargest()和nsmallest()才算高效，不過有些情況下這兩個函數會很方便。

1.4 高效合並有序序列

對於小數據集，將多個有序序列合並到一個新序列很容易。

list(sorted(itertools.chain(*data)))

對於較大的數據集，這個技術可能會占用大量內存。merge()不是對整個合並后的序列排序，而是使用一個堆一次一個元素地生成一個新序列，利用固定大小的內存確定下一個元素。

import heapq
import random

random.seed(2016)

data = []
for i in range(4):
    new_data = list(random.sample(range(1, 101), 5))
    new_data.sort()
    data.append(new_data)

for i, d in enumerate(data):
    print('{}: {}'.format(i, d))

print('\nMerged:')
for i in heapq.merge(*data):
    print(i, end=' ')
print()

由於merge()的實現使用了一個堆，所以它會根據所合並的序列個數消耗內存，而不是根據這些序列中的元素個數。