Pythonで理解する蟻本「2-4 プライオリティキューとヒープ」(p.71)

この記事は「プログラミングコンテストチャレンジブック第2版」（蟻本）の
「2-4 プライオリティキューとヒープ」(p.71)
のコードをPythonで書き直したものとなっています。

ヒープの実装例

MAX_N = 10 ** 8

heap = [0] * MAX_N

sz = 0

def push(x):
    global sz
    
    # 自分のノード番号
    i = sz
    sz += 1
    
    while i > 0:
        # 親のノードの番号
        p = (i - 1) // 2
        
        # もう逆転していないなら抜ける
        if heap[p] <= x:
            break
        
        # 親のノードの数字を下ろして、自分は上に
        heap[i] = heap[p]
        i = p
    
    heap[i] = x

def pop():
    global sz
    
    # 最小値
    ret = heap[0]
    
    # 根に持ってくる値
    sz -= 1
    x = heap[sz]
    
    # 根からおろしていく
    i = 0
    while i * 2 + 1 < sz:
        # 子同士を比較
        a = i * 2 + 1
        b = i * 2 + 2
        if b < sz and heap[b] < heap[a]:
            a = b
        
        # もう逆転していないなら終わり
        if heap[a] >= x:
            break
        
        # 子の数字を持ち上げる
        heap[i] = heap[a]
        i = a
    
    heap[i] = x
    
    return ret

PythonとC++の違い

Pythonではheapqが、C++のpriority_queueに相当します。

しかし完全に同じというわけではなく、C++のpriority_queueが最大値を返す一方で、Pythonのheapqは最小値を返します。

またC++においては、popはコンテナの先頭側の要素を削除するだけで、返り値はありません。

一方でPythonにおいては、heappopはコンテナの先頭側の要素を削除し、その値を返します（返り値がある）。

そのため、蟻本ではスタックの先頭側の要素を出力した後にその値を削除していますが、この記事のコードでは先頭側の要素を削除する際に値を変数に受け取り、その値を出力しています。

（蟻本のコードは「出力→値の削除」の順番になっているが、この記事のコードでは「値の削除→出力」の順番になっている。）

標準ライブラリを使ったコード

コード

import heapq

# 宣言
pque = []

# 要素の追加
heapq.heappush(pque, 3)
heapq.heappush(pque, 5)
heapq.heappush(pque, 1)

# 空になるまでループ
while len(pque) != 0:
    # 最小値の取得および削除
    MIN = heapq.heappop(pque)    # MINに最小値を代入し削除する
    print(MIN)

実行結果

1
3
5

最大値を取り出す目的でheapqを使う

heapqには最大値を取り出すためのメソッドがありません。

そのため、最大値を取り出す目的でheapqを使いたいときは、要素に-1を掛けてヒープに追加し、取り出した後にもう一度-1を掛けるとよいでしょう。

コード

import heapq

# 宣言
pque = []

# 要素の追加
heapq.heappush(pque, -1 * 3)    # 要素を追加する際に-1を掛ける
heapq.heappush(pque, -1 * 5)
heapq.heappush(pque, -1 * 1)

# 空になるまでループ
while len(pque) != 0:
    # 最小値の取得および削除
    MAX = -1 * heapq.heappop(pque)    # 受け取る際にもう一度-1を掛ける
    print(MIN)

実行結果

5
3
1