[Python] 백준 11404번: 플로이드

https://www.acmicpc.net/problem/11404

11404번: 플로이드

 

import sys
input = sys.stdin.readline
INF = 100001


def floydWarshall(n):
    dp = [[INF*100]*(n+1) for _ in range(n+1)]

    # 초기화
    for i in range(1, n+1):
        for j in range(1, n+1):
            if graph[i][j] != INF:
                dp[i][j] = graph[i][j]

    # 플로이드
    for k in range(1, n+1):
        for i in range(1, n+1):
            for j in range(1, n+1):
                if dp[i][k] + dp[k][j] < dp[i][j]:
                    dp[i][j] = dp[i][k] + dp[k][j]

    # 출력
    for i in range(1, n+1):
        for j in range(1, n+1):
            if i == j or dp[i][j] == INF*100:
                print(0, end=" ")
            else:
                print(dp[i][j], end=" ")
        print()


if __name__ == '__main__':
    n = int(input())
    m = int(input())

    graph = [[INF]*(n+1) for _ in range(n+1)]

    for _ in range(m):
        a, b, c = map(int, input().split())
        graph[a][b] = min(graph[a][b], c)

    floydWarshall(n)

 

전형적인 Floyd Warshall 문제이다. 입력을 graph로 받아주고, 이에 대해서 플로이드 와샬 알고리즘을 이용하면 된다.

이 알고리즘은, 모든(every) 정점으로부터 모든(every) 정점까지의 최단거리를 알아낼 수 있으며, 시간복잡도는 O(n^3)이다.

 

floydWarshall() 함수는 아래와 같이 진행된다.

1. 입력받은 그래프를 토대로 dp를 초기화한다. 다만 dp의 초기값에 INF*100을 넣는 것은, dp는 그래프와는 달리 누적된 cost가 저장되기 때문이다.

2. 플로이드 와샬 알고리즘의 메인 로직이다. k는 거쳐가야하는 노드이고, i, j는 탐색하는 방향(i->j)이다. i->k->j일 때 드는 비용이 i->j보다 적으면 dp[i][j]에 dp[i][k] + dp[k][j]를 넣는 식이다. 생각보다 간단하다.

3. 저장되어있는 dp를 출력해준다. 다만, 시작점과 도착점이 같은 경우 또는, dp가 갱신되지 않은 경우(dp가 INF*100인 경우로, 즉 도착할 수 없다는 의미이다.)에는 0을 출력한다.