【Python3】14行でダイクストラ法

昨日は14行で素因数分解を実装したのですが、本日はヒープを用いて14行でダイクストラ法を実装しました。上から14行で。

dijkstra.py

from heapq import heappush, heappop

def dijkstra(vertices_num, edges, src):
  dist = [float('inf')] * vertices_num
  heap = [(0, src)]

  while len(heap):
    d, from_ = heappop(heap)
    if d < dist[from_]:
      dist[from_] = d
      for _, to, cost in filter(lambda e: e[0]==from_, edges):
        heappush(heap, (min(dist[to], dist[from_] + cost), to))
        
  return dist

if __name__ == '__main__':
  answer = dijkstra(
    8, 
  　[
      (0, 1, 5), (0, 2, 4), (0, 3, 1), (1, 4, 2), (1, 0, 5), (2, 0, 4), (2, 3, 2), (2, 4, 5), 
      (2, 5, 6), (3, 0, 1), (3, 2, 2), (4, 1, 2), (4, 2, 5), (4, 6, 1), (4, 7, 3), (5, 2, 6), 
      (5, 7, 2), (6, 4, 1), (6, 7, 4), (7, 4, 3), (7, 5, 2), (7, 6, 4)
    ], 0
  )
  print(answer) # [0, 5, 3, 1, 7, 9, 8, 10]

第1引数は頂点数、第2引数は(from_, to, cost)のリスト、第3引数は開始地点のインデックスになります。返り値は開始地点から各頂点への最短距離を返します。

参考

https://www.youtube.com/watch?v=X1AsMlJdiok
ダイクストラ法に関しては、こちらの動画の説明が非常にわかりやすかったので、そこで扱われている問題を参考にテストケースを作らせて頂きました。赤字は各頂点のインデックスです。

(引用元:上記Youtubeページより)