概要
単一始点最短経路（Single Source Shortest Path; SSSP）の重みを求めるプログラム、通称 Dijkstra（ダイクストラ）法のpythonによる実装を提示します。
アルゴリズム
アルゴリズムについては蟻本をはじめ多くの書籍や記事で詳しく扱われているので、ここでは割愛します。動画で確認したい方は、ヨビノリさんが解説動画を出しているのでリンクを貼っておきます。実際にコードを書く前にアルゴリズムを確認して手を動かすことをお勧めします。ダイクストラのアルゴリズム
実装
ヒープキューを用いて実装します。ヒープキューは要素の追加と最小の要素の取り出しをそれぞれO(log N)で実行できるデータ構造です。pythonではheapqモジュールを用いることができます。
以下にpythonによる実装例を提示します。コード中のadjは二次元配列で、第i要素にノードiと隣接するノードと辺重みの情報をタプルで持っている配列です。
from heapq import heappush, heappop
INF = 10 ** 9
def dijkstra(s, n): # (始点, ノード数)
dist = [INF] * n
hq = [(0, s)] # (distance, node)
dist[s] = 0
seen = [False] * n # ノードが確定済みかどうか
while hq:
v = heappop(hq)[1] # ノードを pop する
seen[v] = True
for to, cost in adj[v]: # ノード v に隣接しているノードに対して
if seen[to] == False and dist[v] + cost < dist[to]:
dist[to] = dist[v] + cost
heappush(hq, (dist[to], to))
return dist
実行イメージ
入力は次のような形式を想定しています。第一行に「ノード数、エッジ数、始点ノード」が与えられ、それ以降の行で「ノード番号、ノード番号、辺重み」が与えられます。
- 入力例
4 5 0 # ノード数, エッジ数, 始点ノード
0 1 1
0 2 4
1 2 2
2 3 1
1 3 5
- 入力の受け取り・隣接リスト
adjの構築：
# ノード数, エッジ数, 始点ノード
v, e, r = map(int, input().split())
# adj[s]: ノード s に隣接する(ノード, 重み)をリストで持つ
adj = [[] for _ in range(v)]
for i in range(e):
s, t, d = map(int, input().split())
adj[s].append((t, d))
- 出力：
>>> print(adj)
[[(1, 1), (2, 4)], [(2, 2), (3, 5)], [(3, 1)], []]
>>> d = dijkstra(r, v)
>>> print(d)
[0, 1, 3, 4]
検証
以下の問題で検証することができます。