ダイクストラ法（Dijkstra's Algorithm）をPythonで実装する

この記事で得られるもの

• ダイクストラ法の手順
• Pythonでの実装例

環境

• Google Colaboratory

背景

大学で離散数学の講義を取っていたのですが、理論を学び、演習として手計算をするという形式でした。

この講義の試験では資料等の持ち込みが許可されていたため、
「アルゴリズム実装して持ち込めば計算しなくて良くない？」
と思い、作成しました。

「アルゴリズムをそのまま実装する」という部分に重きを置いているため、効率や読みやすさは重視していません。

ダイクストラ法の手順

1. 始点の頂点を決定し、その頂点からの最短距離を0と設定する
2. 始点以外の全ての頂点への最短距離を無限大に設定する
3. 未訪問の頂点の中で最短距離が最小の頂点を選ぶ
4. 3で選ばれた頂点から直接接続されている各頂点について、その頂点までの距離（選ばれた頂点までの最短距離 + 選ばれた頂点からその頂点へのエッジの重み）が現在記録されている最短距離よりも小さい場合は、最短距離を更新する
5. 3～4を、未訪問の頂点がなくなるまで繰り返す

実装例

ダイクストラ法

import heapq

def shortest_path(graph, start, end):
    queue = [(0, start, [])]
    seen = set()
    mins = {start: 0}
    while queue:
        (cost, node, path) = heapq.heappop(queue)
        if node not in seen:
            seen.add(node)
            path = path + [node]
            if node == end:
                return cost, path

            for next_node, next_cost in graph[node].items():
                prev_cost = mins.get(next_node, None)
                next_cost += cost
                if prev_cost is None or next_cost < prev_cost:
                    heapq.heappush(queue, (next_cost, next_node, path))
                    mins[next_node] = next_cost

    return float("inf"), []

使用例

以下のグラフを例として考える

# define graph
graph = {
    'v0':{'v1':8,'v2':3, 'v3':1},
    'v1':{'v0':8,'v2':3,'v4':1},
    'v2':{'v0':3,'v1':3,'v4':4},
    'v3':{'v0':1,'v4':2},
    'v4':{'v1':1,'v2':4,'v3':2}
}    

# Create a list of all vertices
vertices = ['v' + str(i) for i in range(0, 5)]

# Iterate over all pairs of vertices
for start in vertices:
    for end in vertices:
        # Skip if start and end are the same
        if start == end:
          cost, path = '-','-'
        else:
          cost, path = shortest_path(graph, start, end)
        print(f"{start} to {end} : {cost}, {path}")

結果

v0 to v0 : -, -
v0 to v1 : 4, ['v0', 'v3', 'v4', 'v1']
v0 to v2 : 3, ['v0', 'v2']
v0 to v3 : 1, ['v0', 'v3']
v0 to v4 : 3, ['v0', 'v3', 'v4']
v1 to v0 : 4, ['v1', 'v4', 'v3', 'v0']
v1 to v1 : -, -
v1 to v2 : 3, ['v1', 'v2']
v1 to v3 : 3, ['v1', 'v4', 'v3']
v1 to v4 : 1, ['v1', 'v4']
v2 to v0 : 3, ['v2', 'v0']
v2 to v1 : 3, ['v2', 'v1']
v2 to v2 : -, -
v2 to v3 : 4, ['v2', 'v0', 'v3']
v2 to v4 : 4, ['v2', 'v4']
v3 to v0 : 1, ['v3', 'v0']
v3 to v1 : 3, ['v3', 'v4', 'v1']
v3 to v2 : 4, ['v3', 'v0', 'v2']
v3 to v3 : -, -
v3 to v4 : 2, ['v3', 'v4']
v4 to v0 : 3, ['v4', 'v3', 'v0']
v4 to v1 : 1, ['v4', 'v1']
v4 to v2 : 4, ['v4', 'v2']
v4 to v3 : 2, ['v4', 'v3']
v4 to v4 : -, -