1238 - 파티 [Dijkstra]

출처 : 백준 1238 파티

더 많은 문제들에 대한 Python3 풀이는 Github Repo에서 찾아볼 수 있다

문제

이번 문제는 왕복으로 최단거리를 구하는 문제이다. 1 -> 2, 2 -> 1로 이동할 때 각각 거리가 다르기 때문에 이 점을 유의해서 풀어야 한다. 처음에 N 개의 집과 M개의 도로 그리고 파티를 하는 장소인 X가 주어진다. 자기 집에서 X까지 갔다가 다시 돌아오는 거리를 왕복거리라고 보고, 가장 왕복거리가 짧은 사람의 index를 리턴해야 한다.

다익스트라 알고리즘은 한 노드에서 다른 노드까지의 최단거리를 구하는 알고리즘이고, 이 문제는 왕복이기 때문에 각 노드에서 X까지 가는 거리를 계산하고, X에서 다른 노드까지 거리를 계산하는 부분을 추가하면 된다. 처음에 N개의 수 만큼 배열을 초기화한 다음에, 다익스트라 알고리즘을 시작 노드를 하나씩 증가시키면서 진행한다. 예를 들어, X가 3일때, 1번을 시작 노드로봐서 알고리즘을 실행한 다음에 나온 결과 중에서 3번째 index(즉, X까지의 거리)에 해당하는 값을 아까 초기화한 배열에 add를 한다.(자신부터 X까지 거리) 그러고 나서 X가 시작 노드가 되서 알고리즘을 실행한 다음에는 그 값만큼 아까 초기화한 배열에 add한다.(X부터 자신까지의 거리)

풀이

# BOJ 1238 - 파티
import sys
from heapq import heappush, heappop
r = sys.stdin.readline
INF = 1e9


def dijkstra(n, x, ro):
    dist = [INF] * n
    dist[x] = 0
    q = []
    heappush(q, [0, x])

    while q:
        cost, pos = heappop(q)

        for p, c in ro[pos]:
            c += cost
            if c < dist[p]:
                dist[p] = c
                heappush(q, [c, p])

    return dist


N, M, X = map(int, r().split())
road = [[] for _ in range(N)]
for _ in range(M):
    u, v, t = map(int, r().split())
    road[u-1].append([v-1, t])

answer = [0] * N
for i in range(N):
    ret = dijkstra(N, i, road)

    if i == X-1:
        for idx, r in enumerate(ret):
            answer[idx] += r
    else:
        answer[i] += ret[X-1]
print(max(answer))