冷めたコーヒー

Weniger, aber besser

Dijkstra 法の隣接行列による実装 in Python

概要

前回の記事において、単一始点最短経路(Single Source Shortest Path; SSSP)の重みを求めるプログラムである Dijkstraダイクストラ)法のpythonによる実装例を紹介しました。そこでは、隣接リストを用いて実装をしましたが、隣接行列を用いても同様に実装することができるので紹介しておきます。

実装

from heapq import heappush, heappop
def dijkstra(s, n): # (始点, ノード数)
    dist = [INF] * n
    hq = [(0, s)] # (distance, node)
    dist[s] = 0
    seen = [False] * n
    while hq:
        v = heappop(hq)[1] # node
        seen[v] = True
        for to in range(n):
            if W[v][to] == INF: continue # 隣接する辺がない場合は処理をスキップ
            if seen[to] == False and dist[v] + W[v][to] < dist[to]:
                dist[to] = dist[v] + W[v][to]
                heappush(hq, (dist[to], to))
    return dist

実行イメージ

次のようなグラフを例に考えてみましょう。

f:id:mirucacule:20200522122701p:plain

  • 入力例:
4 6 1 # ノード数, エッジ数, 始点ノード
0 1 1
0 2 4
2 0 1
1 2 2
3 1 1
3 2 5
  • 入力の受け取り・隣接行列Wの構築:
v, e, r = map(int, input().split()) # ノード数, エッジ数, 始点
INF = 10 ** 4 - 1 # 入力に応じて変更する
W = [[INF] * v for _ in range(v)] # 重み付き隣接行列
for i in range(e):
    s, t, d = map(int, input().split())
    W[s][t] = d
  • 出力:
>>> print(*W, sep='\n') # 隣接行列
[999, 1, 4, 999]
[999, 999, 2, 999]
[1, 999, 999, 999]
[999, 1, 5, 999]
>>> d = dijkstra(r, v)
>>> print(d)
[3, 0, 2, 999]