ワーシャルフロイド法がなぜうまくいくのか、改めて考えてみる

はじめに

グラフの最短経路を見つけるアルゴリズムのひとつに、ワーシャルフロイド法があります。
このアルゴリズムがなぜ正しく動くのか、しっかり考えてみようという記事です。
筆者はアルゴリズムを専攻したわけではないので、もし間違いがあればご指摘頂けると助かります。

実際の問題例

ノード数が$N$の、単純連結無向グラフが与えられる。各ノードには、$0$から$N-1$までの番号が割り振られており、ノード間はいくつかのエッジでつながれている。ノードから別のノードへはエッジを通って移ることができ、その際には各エッジの長さ分の距離を動く必要がある。

始点ノード$i$から終点ノード$j$に到達するための最短移動距離を、すべての$i,j$について求めよ。なお、ノート間を接続するエッジの情報として、２次元配列$dist$が与えられる。$dist[i][j]$の値は以下の通り。

ノード$i$とノード$j$を直接つなぐエッジがある場合には、そのエッジの長さ
ノード$i$とノード$j$を直接つなぐエッジがない場合には、$\infty$
ノード$i=$ノード$j$の場合には、$0$

この問題はワーシャルフロイド法によって解くことができ、pythonコードだと以下のようになります。

for k in range(N):
    for i in range(N):
        for j in range(N):
            dist[i][j] = min(dist[i][j], dist[i][k] + dist[k][j]

非常にシンプルなコードで、実装自体はわかりやすいのですが、なぜこのコードでうまくいくのか考えてみましょうというのが本記事の趣旨です。
まず、この問題にいきなり取り組むのではなく、もっと簡単な問題から考えてみましょう。

制約つきの問題

実際の問題に制約を追加した、以下の問題を考えてみます。

問題０

始点$i$から終点$j$に到達するための最短移動距離を、すべての$i,j$について求めよ。
制約条件：始点$i$から終点$j$への移動では、$i$と$j$以外のノードを一切経由してはいけない。

どぎつい制約があるおかげで、問題としてはとても簡単です。始点から終点への直接移動しか許可されていないのですから、直接の移動距離が格納されている$dist[i][j]$がそのまま答えになります。
では、制約を少しだけ緩和した次の問題ではどうでしょうか。

問題１

始点$i$から終点$j$に到達するための最短移動距離を、すべての$i,j$について求めよ。
制約条件：始点$i$から終点$j$への移動では、$i$と$j$以外のノードとして、ノード$0$しか経由してはいけない。

移動途中にノード$0$のみ経由することが許されるので、$i$から$j$への最短距離は
$i \rightarrow j$
$i \rightarrow 0 \rightarrow j$
のどちらか小さいほうの移動距離ということになります。具体的なコードだと

min(dist[i][j], dist[i][0] + dist[0][j]

になります。全ての始点・終点に対して答えを求めなければいけないので、愚直に２重ループで全探索して$dist$の値を更新してしまいましょう。

# 問題1　ノード0しか経由してはいけない
for i in range(N):
    for j in range(N):
        # i => j と i => 0 => j の移動距離が小さいほうを採用
        dist[i][j] = min(dist[i][j], dist[i][0] + dist[0][j]

ここまで$i=0$や$j=0$の場合を考えていなかったのですが、$dist[0][0]=0$なので、実はこれで問題ありません。

さらに制約をゆるめてみます。

問題2

始点$i$から終点$j$に到達するための最短移動距離を、すべての$i,j$について求めよ。
制約条件：始点$i$から終点$j$への移動では、$i$と$j$以外のノードとして、ノード$0$かノード$1$しか経由してはいけない。

ノード$0$と$1$の経由が許されるので、$i$から$j$への最短距離は
$i \rightarrow j$
$i \rightarrow 0 \rightarrow j$
$i \rightarrow 1 \rightarrow j$
$i \rightarrow 0 \rightarrow 1 \rightarrow j$
$i \rightarrow 1 \rightarrow 0 \rightarrow j$
の５種類の移動方法のうち、移動距離が最も小さいものになります。コードだとこうです。

# 問題2　ノード0と1しか経由してはいけない
for i in range(N):
    for j in range(N):
        # 一番小さいものを採用
        dist[i][j] = min(dist[i][j],                           # ①i => j
                         dist[i][0] + dist[0][j],              # ②i => 0 => j
                         dist[i][1] + dist[1][j],              # ③i => 1 => j
                         dist[i][0] + dist[0][1] + dist[1][j], # ④i => 0 => 1 => j
                         dist[i][1] + dist[1][0] + dist[0][j]) # ⑤i => 1 => 0 => j

なかなか複雑なコードですが、もしも以下のように、問題1の処理が先に書いてあったらどうなるでしょうか。

# 問題1　ノード0しか経由してはいけない
for i in range(N):
    for j in range(N):
        # i => j と i => 0 => j の移動距離が小さいほうを採用
        dist[i][j] = min(dist[i][j], dist[i][0] + dist[0][j]

# 問題2　ノード0と1しか経由してはいけない
for i in range(N):
    for j in range(N):
        # 一番小さいものを採用
        dist[i][j] = min(dist[i][j],                           # ①i => j
                         dist[i][0] + dist[0][j],              # ②i => 0 => j
                         dist[i][1] + dist[1][j],              # ③i => 1 => j
                         dist[i][0] + dist[0][1] + dist[1][j], # ④i => 0 => 1 => j
                         dist[i][1] + dist[1][0] + dist[0][j]) # ⑤i => 1 => 0 => j

この場合、問題2の処理でminをとる5個ののうち、

dist[i][j]                # ①i => j
dist[i][0] + dist[0][j]   # ②i => 0 => j

については、minをとるまでもなく、

dist[i][j]

にまとめてしまってもOKです。なぜなら、直接移動するのとノード$0$を経由するのとで、移動距離の小さいほうの値が既に$dist[i][j]$に格納されているからです。残りも見てみましょう。

dist[i][1] + dist[1][j]              # ③i => 1 => j
dist[i][0] + dist[0][1] + dist[1][j] # ④i => 0 => 1 => j
dist[i][1] + dist[1][0] + dist[0][j] # ⑤i => 1 => 0 => j

すると、

dist[i][0] + dist[0][1]
dist[1][0] + dist[0][j]

の部分を

dist[i][1]
dist[1][j]

に置き換えてしまって問題ないことが分かります。つまり③④⑤は

dist[i][1] + dist[1][j]

にまとめてしまってOKです。結局、minの中身は5個から2個に減り、シンプルなコードになります。

# 問題1　ノード0しか経由してはいけない
for i in range(N):
    for j in range(N):
        dist[i][j] = min(dist[i][j], dist[i][0] + dist[0][j]
# 問題2　ノード0と1しか経由してはいけない
for i in range(N):
    for j in range(N):
        dist[i][j] = min(dist[i][j], dist[i][1] + dist[1][j]

似たような処理が並んでいるので、for文でひとまとめにしてしまいましょう。

# 問題2　ノード0と1しか経由してはいけない
for k in range(2):
    for i in range(N):
        for j in range(N):
            dist[i][j] = min(dist[i][j], dist[i][k] + dist[k][j]
    # k=0の時点では、問題1の答えがdistに格納されている

同じように、どんどん制約を緩めて問題3,4...と考えていくと、実は最外ループのk上限を変えるだけで対応できてしまうことが分かります。元の問題は、$N$個すべてのノードを途中経由しても良いという「問題N」に対応するので、

# 問題N　0からN-1までの全ノードを途中経由しても良い
for k in range(N):
    for i in range(N):
        for j in range(N):
            dist[i][j] = min(dist[i][j], dist[i][k] + dist[k][j]
    # この時点では、途中経由してもよいノードを0からkまでに制限した、問題k+1の答えがdistに格納されている

という、ワーシャルフロイド法として最初に紹介したコードでOKなわけです。

おわりに

この記事を読んで頂けた方には、ぜひAtcoderのBeginnerContest208-D問題も解いてみることを強くおすすめします。ワーシャルフロイド法をただ知っているというだけでなく、なぜうまくいくのかを理解していて初めて解ける問題となっています。
https://atcoder.jp/contests/abc208/tasks/abc208_d