[Python] ABC012 D ワーシャルフロイド

問題

D – バスと避けられない運命

ワーシャルフロイド

ワーシャル–フロイド法（英: Warshall–Floyd Algorithm）は、重み付き有向グラフの全ペアの最短経路問題を多項式時間で解くアルゴリズムである。

出典: フリー百科事典『ウィキペディア（Wikipedia）』

以下の動画がアルゴリズムの概要がつかみやすいです。

回答

工夫すると通るようですが、Python では TLEでした。

import sys

N, M = map(int, input().split())

# 距離行列の生成
dist_matrix = [[sys.maxsize] * (N) for _ in range(N)] 
for i in range(N):
    dist_matrix[i][i] = 0
for i in range(M):
    a, b, t = map(int, input().split())
    dist_matrix[a-1][b-1] = t
    dist_matrix[b-1][a-1] = t

#ワーシャルフロイド
for k in range(N):
    for i in range(N):
        for j in range(N):
            dist_matrix[i][j] = min(dist_matrix[i][j], dist_matrix[i][k]+dist_matrix[k][j])

ans = sys.maxsize
for i in range(N):
    # スタートからの最長距離
    max_dist = -1
    for j in range(N):
        dist = dist_matrix[i][j]
        max_dist = max(max_dist, dist)
    # 最長距離が最も短くなる
    ans = min(ans, max_dist)

print(ans)
    

scipy.sparse.csgraph.floyd_warshall を使ってみます。

from scipy.sparse import csr_matrix
from scipy.sparse.csgraph import floyd_warshall

N, M = map(int, input().split())

# 距離行列の生成
dist_matrix = [[0] * (N) for _ in range(N)] 
for i in range(N):
    dist_matrix[i][i] = 0
for i in range(M):
    a, b, t = map(int, input().split())
    dist_matrix[a-1][b-1] = t
    dist_matrix[b-1][a-1] = t

graph = csr_matrix(dist_matrix)
# ワーシャルフロイド
dist_matrix = floyd_warshall(csgraph=graph, directed=False)

ans = min([max(row) for row in dist_matrix[:]])
print(ans)