[BOJ] 1238 파티 - 다익스트라(Dijkstra) 풀이

문제

N개의 숫자로 구분된 각각의 마을에 한 명의 학생이 살고 있다.

어느 날 이 N명의 학생이 X (1 ≤ X ≤ N)번 마을에 모여서 파티를 벌이기로 했다. 이 마을 사이에는 총 M개의 단방향 도로들이 있고 i번째 길을 지나는데 Ti(1 ≤ Ti ≤ 100)의 시간을 소비한다.

각각의 학생들은 파티에 참석하기 위해 걸어가서 다시 그들의 마을로 돌아와야 한다. 하지만 이 학생들은 워낙 게을러서 최단 시간에 오고 가기를 원한다.

이 도로들은 단방향이기 때문에 아마 그들이 오고 가는 길이 다를지도 모른다. N명의 학생들 중 오고 가는데 가장 많은 시간을 소비하는 학생은 누구일지 구하여라.

풀이 (Python)

import heapq

n, m, x = map(int, input().split())
adj = [[] for _ in range(n+1)]

for _ in range(m):
    start, end, t = map(int, input().split())
    adj[start].append((end, t))


def dijkstra(start):
    INF = float('inf')

    q = []
    heapq.heappush(q, (0, start))
    dp = [INF]*(n+1)
    dp[start] = 0

    while q:
        cost, now = heapq.heappop(q)
        if dp[now] < cost:
            continue

        for next, c in adj[now]:
            new_cost = cost+c
            if new_cost < dp[next]:
                dp[next] = new_cost

            heapq.heappush(q, (new_cost, next))

    return dp

ans = 0
x_to_i = dijkstra(x)

for i in range(1, n+1):
    if i == x:
        continue

    temp_dp = dijkstra(i)
    time = temp_dp[x] + x_to_i[i]
    if ans < time:
        ans = time

print(ans)

해설

시간 초과와 메모리 초과를 모두 겪은 문제였습니다.
BFS로 접근해서 방문 배열을 들고 다니는 식으로 풀었다가 메모리 초과를 겪었고,
다익스트라 알고리즘을 사용해서 통과했습니다.
다익스트라 알고리즘을 비롯해 최단경로에 대한 이론은 여기에서 보실 수 있습니다.

문제의 조건에 따르면 i번 마을에서 출발해서 x번 마을을 경유해 다시 i번 마을로 돌아와야 합니다.
즉 x를 경유하는 i->i 최단거리 사이클을 찾아야 하는 문제입니다.

다익스트라 알고리즘은 하나의 기준 노드로부터 다른 모든 노드로부터의 최단거리를 찾아주는 알고리즘입니다.
그래서 이 문제를 다익스트라 알고리즘으로 접근해보면,
x를 경유하는 i->i 사이클은 i->x 최단거리와 x->i 최단거리의 합이 될 것입니다.

따라서 다익스트라 알고리즘을 구현한 함수를 이용해 최단거리를 계산한 dp 배열을 구하고,
i->x 최단거리와 x->i 최단거리의 합 중 최대값을 구해주면 됩니다.
이때 x->x인 경우를 제외하는 것도 잊지 마세요.

x->i 최단거리는 고정된 값이기 때문에 for문에 들어가기 전에 미리 한 번만 구해놓고 시작해야 시간초과를 피할 수 있습니다.
저는 이 부분에 최적화를 안 해서 시간초과가 났던 것 같습니다.