1719. 택배 (Python)

1719번: 택배

 

풀이

 플로이드 워셜 알고리즘을 사용하여 해결할 수 있다. 기존의 플로이드 워셜 알고리즘에서 최단거리를 기록하는 graph 배열에 첫 번째로 방문하는 집하장 정보를 추가한다. i에서 j로 가는 경로가 k 집하장을 거쳐서 가는 경로의 길이로 갱신될 때 graph[i][k]에 기록되어 있는 첫 번째로 방문하는 집하장 정보를 graph[i][j]에 기록해주면 된다.

 

코드

import sys
input = sys.stdin.readline
INF = float('inf')


if __name__ == '__main__':
    N, M = map(int, input().split())
    # graph[i][j] -> [i에서 j로 가는 최단경로의 길이, i에서 j로 갈 때 처음으로 가는 집하장]
    # 길이는 INF로, 처음 가는 집하장 정보는 자기 자신으로 초기화
    graph = [[[INF, i] for i in range(N + 1)] for _ in range(N + 1)]

    for _ in range(M):  # 경로 입력
        a, b, time = map(int, input().split())
        graph[a][b][0] = time
        graph[b][a][0] = time

    # 자기 자신으로 가는 경로의 길이는 0으로 초기화
    for i in range(1, N + 1):
        graph[i][i][0] = 0

    # 플로이드 워셜 알고리즘
    for k in range(1, N+1):
        for i in range(1, N+1):
            for j in range(1, N+1):
                # k를 거쳐가는 경로로 갱신될 때, 처음 거치는 집하장 정보를 갱신해준다.
                if graph[i][j][0] > graph[i][k][0] + graph[k][j][0]:
                    graph[i][j][0] = graph[i][k][0] + graph[k][j][0]
                    graph[i][j][1] = graph[i][k][1]

    for i in range(1, N+1):
        for j in range(1, N+1):
            if i == j:
                print('-', end=' ')
            else:
                print(graph[i][j][1], end=' ')
        print()