【競技プログラミング:リハビリ】ワーシャルフロイド

既存投稿一覧ページへのリンク

一覧ページ

関連する問題

ワーシャルフロイド

sample_code.py

def warshall_floyd(edges):
    # ノードの最大番号を取得
    max_node = max(max(e[0], e[1]) for e in edges)
    n = max_node + 1  # ノード数
    
    # 距離行列の初期化
    INF = float('inf')
    dist = [[INF] * n for _ in range(n)]
    
    # 自分自身への距離は0
    for i in range(n):
        dist[i][i] = 0
    
    # エッジ情報を反映
    for src, dst, val in edges:
        if dist[src][dst] > val:  # 重複エッジがある場合は最小値を採用
            dist[src][dst] = val
    
    # ワーシャルフロイド法のメイン処理
    for k in range(n):
        for i in range(n):
            for j in range(n):
                if dist[i][k] != INF and dist[k][j] != INF:
                    dist[i][j] = min(dist[i][j], dist[i][k] + dist[k][j])
    
    return dist

# 使用例
if __name__ == "__main__":
    edges = [
        [0, 1, 5],
        [1, 2, 3],
        [0, 2, 9],
        [1, 4, 3],
        [3, 4, 1],
    ]
    result = warshall_floyd(edges)
    
    # 結果表示
    print("最短距離行列:")
    for row in result:
        print(" ".join(f"{x:2}" if x != float('inf') else "∞" for x in row))