素敵なデータ構造：ヒープ by Python

データ構造の説明

ヒープとは、
・二分木を用いるデータ構造で、
・木の深さh-1以下の部分については完全二分木を形成し、
・各頂点と親頂点との間に親頂点≧子頂点の関係が成立し、
・値xの検索には適さないですが、最大値の取得が得意
なようです。

heap.py

class Heap():
    heap = []

    # 値の追加
    # 一番最後の葉に追加したのち、「自分の親と比較して逆転していれば入れ替える」を繰り返す
    def push(self, x):
        self.heap.append(x)
        i = len(self.heap)-1
        while i > 0:
            p = (i-1)//2
            if self.heap[p] >= x:
                break
            self.heap[i] = self.heap[p]
            i = p
        self.heap[i] = x
        return self.heap

    # 最大値の取得
    def top(self):
        if self.heap != "":
            return self.heap[0]
        else:
            return -1

    # 最大値の削除
    # 一番最後の葉を根と入れ替えたのち、「自分の子と比較して逆転していれば入れ替える」を繰り返す
    def pop(self):
        if self.heap == "":
            return
        x = self.heap[-1]
        self.heap.pop()
        i = 0
        while i*2+1 < len(self.heap):
            child1 = i*2+1
            child2 = i*2+2
            if child2 < len(self.heap) and self.heap[child2] > self.heap[child1]:
                child1 = child2
            if self.heap[child1] <= x:
                break
            self.heap[i] = self.heap[child1]
            i = child1
        self.heap[i] = x
        return self.heap

h = Heap()
h.push(3)
h.push(43)
h.push(15)
h.push(29)
h.pop()
h.push(732)
h.push(88)
h.push(14)
h.push(51)
h.pop()

# 出力 [88, 29, 51, 3, 15, 14]

アルゴリズムへの適用

ダイクストラ法

Dijkstra-heap.py

import heapq

# アルゴリズム定義
# ヒープを用いたダイクストラ法
def Dijkstraheap(G, dist, s):
    used = [False for i in range(N)]
    dist[s] = 0    
    que = []
    heapq.heappush(que, (dist[s], s))

    while que:
        v = que[0][1]
        d = que[0][0]
        heapq.heappop(que)

        if d > dist[v]:
            continue

        for e in range(len(G[v])):
            if dist[G[v][e][0]] > dist[v] + G[v][e][1]:
                dist[G[v][e][0]] = dist[v] + G[v][e][1]
                heapq.heappush(que, (dist[G[v][e][0]], G[v][e][0]))

# 入力受取
# 頂点数、辺数、始点
N, M, s = map(int, input().split())

# グラフの生成
# 重み付き単方向グラフ
G = [[] for i in range(N)]
for i in range(M):
    a, b, w = map(int, input().split())
    G[a].append([b, w])

# ダイクストラ法の結果を格納するリストの生成 
# 初期値として無限大の代わりに10000を格納
dist = [10000 for i in range(N)]    

# アルゴリズム実行
# ヒープを用いたダイクストラ法
Dijkstraheap(G, dist, s)

# 結果出力
# sから各頂点への最短距離を出力
for i in range(N):
    print(i, ": ", dist[i])

# 入力
# 6 9 0
# 0 1 3
# 0 2 5
# 1 2 4
# 1 3 12
# 2 3 9
# 2 4 4
# 3 5 2
# 4 3 7
# 4 5 8

# 出力
# 0 :  0
# 1 :  3
# 2 :  5
# 3 :  14
# 4 :  9
# 5 :  16