AtCoder ABC208 D - Shortest Path Queries 2

ワーシャル・フロイド法の理解を問う問題です。緑diffです。

問題

$G$ を $N$ 頂点 $M$ 辺の有向グラフ、 $f(s,t,k)$ を $s$ から $t$ へ頂点番号 $k$ 以下の頂点のみ経由したときの最小距離とする。 $\displaystyle \sum_{s=1}^N \sum_{t=1}^N \sum_{k=1}^N f(s,t,k)$ を求めよ。

制約

$1 \le N \le 400$
$M \le N(N-1)$

考察

全点間の距離なので、計算量 $O(N ^3)$ のワーシャル・フロイド法が思いつきます。しかし、 $k=1,\cdots,N$ に対してワーシャル・フロイド法を用いると、全体で $O(N ^4)$ になり、間に合いません。

ここで、ワーシャル・フロイド法のコード

for k in range(N):
    for i in range(N):
        for j in range(N):
            if cost[i][k] != INF and cost[k][j] != INF:
                cost[i][j] = min(cost[i][j], cost[i][k] + cost[k][j])

を眺めると、一番外側の $k$ のループがまさに、頂点 $k$ を付加して全点間の距離を更新する処理になっていることが分かります。

それで $\text{cost}_{i, j}$ を、

$\text{cost}_{u, u} = 0 (u=1,\cdots,N)$
$\text{cost}_{A_i, B_i} = C_i$
$\text{cost}_{\text{other}, \text{other}} = \infty$

で初期化し、 $k$ の順に処理します。 $k$ の時の全点間の距離の和は差分処理で求めています。

実装例

INF = 10**15

N, M = map(int, input().split())
cost = [[INF] * N for _ in range(N)]
for i in range(N):
    cost[i][i] = 0

for _ in range(M):
    a, b, c = map(int, input().split())
    cost[a-1][b-1] = c

v = 0
for a in range(N):
    for b in range(N):
        v += (cost[a][b] if cost[a][b] < INF else 0)

ans = 0
for k in range(N):
    for a in range(N):
        for b in range(N):
            d = cost[a][k] + cost[k][b]
            if d < cost[a][b]:
                v += -(cost[a][b] if cost[a][b] < INF else 0) + d
                cost[a][b] = d

    ans += v

print(ans)