백준 1238번 : 파티 in Python

1238번: 파티

 

www.acmicpc.net

이 문제는 N개의 숫자로 구분된 각각의 마을에 살고 있는 N명의 학생이 X번 마을에 모여서 파티를 벌인 뒤 각자의 마을로 최단 시간에 돌아갈 때, 오고 가는데 가장 오래 걸린 학생을 찾는 문제이다. 단, 주어지는 도로들은 모두 단방향이다.

우선 X번 마을에서 출발해 각자의 마을로 돌아가는데 걸리는 최단 시간을 구해보자. 이는 각 마을을 정점으로 하고 각 도로를 방향이 있는 간선으로 하는 그래프로 해석하면 각 마을까지의 최단 시간은 X번 마을을 시작점으로 하는 dijkstra 알고리즘을 활용해 구할 수 있다. 

하지만 각자의 마을에서 X번 마을에 도착하는 최단 시간을 구하기 위해서 위 dijkstra 알고리즘을 바로 적용할 경우 각 마을을 시작점으로 하는 dijkstra 알고리즘을 실행, 즉 dijkstra 알고리즘을 N번 실행한 뒤 각 결과에서 목적지가 X번 마을일 때의 최단 시간만을 골라 찾아야 한다. 이는 매우 비효율적인 알고리즘이다. 왜냐하면 임의의 한 마을을 시작점으로 하는 dijkstra 알고리즘을 통해 구한 각 마을까지의 최단 시간 중 X번 마을이 목적지일 때를 제외하고 모두 버려지기 때문이다. 또한 이를 실행하는데 걸리는 시간은 O(N (N+M) logN) = O(N^3 logN) (N : 정점의 개수, M : 간선의 개수로 O(N^2)이라 할 수 있다.)이기 때문에 시간적인 면에서도 비효율적이라 할 수 있다.

그러므로 더 효율적인 방법, 즉 각 마을을 시작점으로 하고 X번 마을만 목적지로 하는 알고리즘을 구현할 필요가 있다. 이를 위해 발상의 전환이 필요하다. 만약 A번 마을에서 X번 마을까지의 최단 경로를 구했다고 하자. 그러면 그 반대 경로는 X번 마을에서 A번 마을까지의 경로가 될 것이고, 경로 내 각 간선은 기존의 그래프 내 간선을 방향만 반대로 한 상태가 될 것이다. 즉, 각 마을을 시작점으로 하고 X번 마을만 목적지로 하는 최단 경로들을 구하기 위해, 그래프의 모든 간선을 반대로 뒤집고 반대로 X번 마을이 시작점이 되고 각 마을이 목적지가 되도록 바꾼다. 그러면 각 마을까지 도달하는 최단 경로를 구하는 문제는 이전의 dijkstra 알고리즘을 활용해 쉽게 해결할 수 있고, 또한 필요한 정보만 얻을 뿐 아니라 계산에 걸리는 시간도 줄여 효율성도 얻을 수 있다. 

따라서 위 내용을 적용하기 위해 그래프를 저장하는 배열을 두 개 만들어 하나는 주어진 간선을 그대로 저장하고, 다른 하나는 주어진 간선을 방향을 바꿔 저장한다. 그리고 각 배열에 대해 X번 마을을 시작점으로 하는 dijkstra 알고리즘을 적용해 각 마을까지의 최단 시간을 얻는다. 마지막으로 각 마을까지 걸리는 시간의 합의 최댓값을 구하면 된다.

이를 코드로 구현하면 다음과 같다.

import sys
import heapq
N, M, X = map(int, input().split())
graph = [[] for _ in range(N+1)]
# 주어진 각 간선의 방향을 바꿔 그래프를 저장한다.
graph_rev = [[] for _ in range(N+1)]
for _ in range(M):
    start, end, cost = map(int, input().split())
    graph[start].append([end, cost])
    graph_rev[end].append([start, cost])
# dijkstra 알고리즘을 구현하는 함수
def dijkstra(_start, _graph):
    dijk = [sys.maxsize for _ in range(N+1)]
    visited = [False for _ in range(N+1)]
    dijk[_start] = 0
    q = []
    heapq.heappush(q, [0, _start])
    while q:
        d, u = heapq.heappop(q)
        if visited[u]: continue
        visited[u] = True
        for v, e in _graph[u]:
            if dijk[v] > d + e:
                dijk[v] = d + e
                heapq.heappush(q, [dijk[v], v])
    return dijk
# 각 마을에서 X번 마을까지 가는데 걸리는 최단 시간을 구한다.
toX = dijkstra(X, graph_rev)
# X번 마을에서 각 마을까지 가는데 걸리는 최단 시간을 구한다.
fromX = dijkstra(X, graph)
res = 0
for i in range(1, N+1):
    if res < toX[i] + fromX[i]:
        res = toX[i] + fromX[i]
print(res)