グラフ系アルゴリズム:ダイクストラ法

今回は競技プログラミングの記事です。

グラフのある1点から、他のすべての点への最短距離を求めらえれるアルゴリズムである、ダイクストラ法(Dijkstra's algorithm)のご紹介です。ご存じな方も多いと思いますが、ダイクストラ法が適用できるグラフは、辺のコストがすべて非負であるものに限ります。

さっそくスクリプトをどうぞ。

import math
import heapq

class Dijkstra:
    def __init__(self, graph):
        self.n = len(graph)
        self.graph = graph
    
    def solve(self, start):
        dp = [[math.inf, False] for _ in range(self.n)]
        remainings = self.n
        dp[start][0] = 0
        q = [[0, start]]
        while q and remainings:
            cost_s2i, i = heapq.heappop(q)
            if dp[i][1]:
                continue
            dp[i][1] = True
            remainings -= 1
            for j, cost_i2j in self.graph[i]:
                if (not dp[j][1]) and cost_s2i + cost_i2j < dp[j][0]:
                    dp[j][0] = cost_s2i + cost_i2j
                    heapq.heappush(q, [dp[j][0], j])
        return [dpi[0] for dpi in dp]

オブジェクト指向な感じのスクリプトです。使い方としてはこんな感じになります。

# graph は隣接リストで用意いただく想定です。
graph = [[[...]]]

d = Dijkstra(graph)
# 距離を測りたい始点のインデックスを渡してください。
x = 0
d.solve(x)
# => [d0, d1, ... dn-1]
# x から各点への最短距離がリストで手に入ります。
# なお、常に dx = 0 です。

深さ優先探索や幅優先探索の場合は、探索のために deque を使いましたがダイクストラ法では heapq を使います。ダイクストラ法では今まで探索を終えた点からたどり着ける点の中で、もっとも近い距離にいる点から順番に探索をしていきます。

heapq 自体はこちらの Python の公式ドキュメントをお読みください。概要としては、抜粋ですが以下をどうぞ。

ヒープとは、全ての親ノードの値が、その全ての子の値以下であるようなバイナリツリーです。この実装は、全ての k に対して、ゼロから要素を数えていった際に、heap[k] <= heap[2*k+1] かつ heap[k] <= heap[2*k+2] となる配列を使っています。比較のために、存在しない要素は無限大として扱われます。ヒープの興味深い性質は、最小の要素が常にルート、つまり heap[0] になることです。

Python の heapq は、最小の要素の取得(heappop)と最大の要素の取得(heappush)を、heapq の要素数 N として、O(log(N)) で実行することができます。

さて、ダイクストラ法に戻りまして。heapq を使って、常に heapq の中にある(=今まで探索を終えた点からたどり着ける点の中にある)点のうち、始点からの距離が最小のものから順番に探索するようにすると、探索が完了した点は最短距離が確定します。

というのも、始点から未探索のノードまでの最短距離は、

  • heapq の中にすでに積まれているもの
  • heapq の中に積まれているものを探索済みにすることで、新たにたどり着けるようになるノードとの最短距離

のどちらかとなるので、すべての辺のコストが非負の場合、必ず heapq の中にある値の最小値以上となるためです。

ダイクストラ法の計算量は、辺の数を  E、ノードの数を  V とすると、heapq の計算量から  O( (E + V) log(V) ) となるそうです。

ただ、私のスクリプトだと  O( (E + V) log(E+V) ) なような気も。私のスクリプトか、私の計算量の見積もりのどちらかが間違えているのかもしれないです。

みなさまお使いの際は、改善できそうでしたら適宜処理を改善して使っていただければと思います。*1

*1:世界にはフィボナッチキューなるものもあるらしく、heapq のかわりにそちらを使うと、さらに計算量が改善するそうです。