最近早水先生の講義を見ているのですが, そこでダイクストラ法の説明があり, 実装したくなったので実装した.
import heapq INF = float("inf") # 各駅間の道を (行き先, 重み) という要素の隣接リストの形で保持する. graph = [ [(1, 16), (2, 8), (3, 15)], [(0, 16), (2, 13), (4, 27), (6, 78)], [(0, 8), (1, 13), (3, 21), (4, 10), (5, 17)], [(0, 15), (2, 21), (5, 33)], [(2, 10), (1, 27), (6, 15)], [(2, 17), (3, 33), (6, 52)], [(1, 78), (4, 15), (5, 52)], ] # 英語バージョン """ stations = [ "Tokyo", "Shibuya", "Shinagawa", "Sinjuku", "Shinyokohama", "Yokohama", "Odawara", ] """ # 日本語バージョン stations = [ "東京", "渋谷", "品川", "新宿", "新横浜", "横浜", "小田原", ] def dijkstra(s, n): # (始点, ノード数) dist = [INF] * n hq = [(0, s)] # (distance, node) heapq.heapify(hq) dist[s] = 0 prevs = [-1] * n # 前の駅(経路復元用) seen = [False] * n # ノードが確定済みかどうか while hq: cost, v = heapq.heappop(hq) # ノードをpopする seen[v] = True for to, cost in graph[v]: # ノードvに隣接しているノードに対して if seen[to] == False and dist[v] + cost < dist[to]: print(f"{stations[to]}のコストを更新! {dist[to]} -> {dist[v]+cost}") dist[to] = dist[v] + cost prevs[to] = v heapq.heappush(hq, (dist[to], to)) return dist, prevs # (出発, 到着) start, end = input().split() if start not in set(stations) or end not in set(stations): print("指定された駅は無いよ!") exit(1) dist, prevs = dijkstra(stations.index(start), len(stations)) print(f"{start}から{end}の所要時間は {dist[stations.index(end)]}") # 経路復元 def restore_path(prevs, v): if prevs[v] != -1: restore_path(prevs, prevs[v]) print(f"{stations[v]}", end="->" if v != stations.index(end) else "\n") restore_path(prevs, stations.index(end))
実行すると
$ python3 dijkstra.py 東京 小田原 渋谷のコストを更新! inf -> 16 品川のコストを更新! inf -> 8 新宿のコストを更新! inf -> 15 新横浜のコストを更新! inf -> 18 横浜のコストを更新! inf -> 25 小田原のコストを更新! inf -> 94 小田原のコストを更新! 94 -> 33 東京から小田原の所要時間は 33 東京->品川->新横浜->小田原
とこんな感じで講義動画と同じような経路更新が出来ている.
他の経路も試してみる
$ python3 dijkstra.py 渋谷 新宿 東京のコストを更新! inf -> 16 品川のコストを更新! inf -> 13 新横浜のコストを更新! inf -> 27 小田原のコストを更新! inf -> 78 新宿のコストを更新! inf -> 34 新横浜のコストを更新! 27 -> 23 横浜のコストを更新! inf -> 30 新宿のコストを更新! 34 -> 31 小田原のコストを更新! 78 -> 38 渋谷から新宿の所要時間は 31 渋谷->東京->新宿 $ python3 dijkstra.py 渋谷 横浜 東京のコストを更新! inf -> 16 品川のコストを更新! inf -> 13 新横浜のコストを更新! inf -> 27 小田原のコストを更新! inf -> 78 新宿のコストを更新! inf -> 34 新横浜のコストを更新! 27 -> 23 横浜のコストを更新! inf -> 30 新宿のコストを更新! 34 -> 31 小田原のコストを更新! 78 -> 38 渋谷から横浜の所要時間は 30 渋谷->品川->横浜 $ python3 dijkstra.py 小田原 東京 渋谷のコストを更新! inf -> 78 新横浜のコストを更新! inf -> 15 横浜のコストを更新! inf -> 52 品川のコストを更新! inf -> 25 渋谷のコストを更新! 78 -> 42 東京のコストを更新! inf -> 33 渋谷のコストを更新! 42 -> 38 新宿のコストを更新! inf -> 46 横浜のコストを更新! 52 -> 42 小田原から東京の所要時間は 33 小田原->新横浜->品川->東京
正しく求められているっぽい.
感想
早水先生の講義わかり易すぎてビビる.
乗り換えアプリとかも基本的にはこういう感じで作られているのかなぁと思った.
最小値を取り出すところはPriotiryQueueを使うと最小値を計算時間(O(log N))で取得出来て良い感じ.
