ダイクストラ法をPythonで実装して関数化する

hirohirohirohiros.hatenablog.com

概要

　ダイクストラ法はグラフの最短路を求めるアルゴリズムです．グラフの2頂点が与えられたとき，その2頂点を始点終点としたときのルートで最もコストが小さいものを最短路と言います．ダイクストラ法は始点を固定したとき，他の全ての頂点への最短路を求めることが出来ます．

アルゴリズムの内容

　コストがマイナスの辺が無いときを考えます．最短路が確定しているノードから最もコストの低い辺のノードに移動した時，移動した先のノードの最短路は，移動する前のノードの最短路＋辺のコストで確定します．これはコストがマイナスの辺が無いからです．最短路が確定したノードはそれ以下のコストが存在しないということなので今後探索する必要はありません．

　手順としては

1. 始点のノードを最短路0として確定する．
2. 最短路が確定しているノードの中から，始点からの最短路が最も小さいノードを取り出す．
3. 取り出したノードに隣接するノードでまだ最短路が確定していないノードがあれば最短路を計算し確定する．
4. 始点から移動可能な全てのノードについて, 最短路が確定するまで2,3を繰り返す．

　2の最短路が最も小さいノードを見つける所をそのまま実装するとO(V)となります(V=ノードの数)．しかし，やりたいことは最小値の取り出しであるのでヒープを使えばO(log V)で出来ます．3の隣接するノードの取り出しは辺の数(=E)だけ行われるのでO(E)かかります．よって全体の計算量はO(E log V)となります．

pythonコード

　以上を実装しました．競プロで使う場合，大抵始点は0になると思われますが始点をstartで選択出来るようにし，デフォルトでは0にするようにしました．

import heapq

def dijkstra(costs, start=0):
    """ダイクストラ法

    startノードを始点とし, 他のノードを終点として移動したとき, 通る辺のコストの和の最小値
    負の辺が存在する時使えない
    計算量: O(E log V)

    Args:
        costs (list): グラフの隣接リスト. 2次元リスト. len(costs)=グラフのノードの総数.
                      costs[n] = n番目のノードから隣接する[(x, x番目のノードに移動するコスト), (y, y番目のノードに移動するコスト),...]

        start (int): 始点のノード. デフォルトでは0

    Returns:
        dist (list): startノードから[0番目のノードに移動するコストの最小値, 1番目のノードに移動するコストの最小値, ...]
                     dist[start] = 0
    """
    INF = 10**18
    dist = [INF] * len(costs)
    dist[start] = 0
    que = [(0, start)] #(コスト, ノード)
    visited = [False] * len(costs)

    while que:
        _, q = heapq.heappop(que) #ノードの取り出し
        visited[q] = True

        #ノードqに隣接するノードに対して最短距離を確定する
        for node, cost in costs[q]: 
            if not visited[node] and dist[q] + cost < dist[node]:
                dist[node] = dist[q] + cost
                heapq.heappush(que, (dist[node], node))

    return dist