【ABC208】 D - Shortest Path Queries 2

問題リンク

ワーシャルフロイド法は名前も知っていたし、使う問題も解いたことがあったが、アルゴリズムの理解が足りなかったため、この問題はコンテスト中に解くことができなかった。

解法

ワーシャルフロイド法は、
dp[k][s][t]: k以下の頂点を中継してよいときのsからtまでの最短経路
といったDPをするアルゴリズムで、問題文と非常に酷似している。

DP配列の初期状態は、
・s = tのとき: dp[0][s][t] = 0
・sからtへ到達不可能のとき: dp[0][s][t] = INF
・sからtへ辺が存在するとき: dp[0][s][t] = 辺の長さ
となる。

sからtへの中継地点として、
・kを使用するとき: dp[k][s][k] + dp[k][k][t]
・kを使用しないとき: dp[k][s][t]
の二つのパターンの最小値をとることで値を更新していく。

ワーシャルフロイド法についての解説は下の記事がわかりやすい。
素人によるワーシャルフロイド法

https://blog.hamayanhamayan.com/entry/2021/07/05/013220

---

実装

dpは2次元でもいいが、定義としては3次元で説明されていることも多いので、理解を深めるためにも3次元で実装してみた。
最初、遷移の部分を、chmin(dp[k+1][s][t], dp[k][s][k] + dp[k][k][t])のように書いてしまったが、これではkを通らない場合を選択できないので間違い。

void solve(long long N, long long M, std::vector<long long> A, std::vector<long long> B, std::vector<long long> C) {
    
    vector<vector<vint>> dp(N+1, vector(N, vint(N, INF)));

    // 初期状態の構築
    REP (i, N) dp[0][i][i] = 0;
    REP (i, M) {
        dp[0][A[i]-1][B[i]-1] = C[i];
    }

    // ワーシャルフロイド
    ll ans = 0;
    REP (k, N) {
        REP (s, N) {
            REP (t, N) {
                dp[k+1][s][t] = min(dp[k][s][t], dp[k][s][k] + dp[k][k][t]);
                if (dp[k+1][s][t] != INF) ans += dp[k+1][s][t];
            }
        }
    }

    c(ans)
}