[알고리즘] 다익스트라 최단 경로 예제

이것이 취업을 위한 코딩 테스트다. with 파이썬 책을 참고하여 정리한 내용입니다.

📖 전보

어떤 나라에는 N개의 도시가 있다. 각 도시에서 다른 도시로 전보를 통해 메시지를 전송할 수 있다. 하지만 X라는 도시에서 Y라는 도시로 전보를 보내고자 한다면, 도시 X에서 Y로 향하는 통로가 있어야 한다. 예를 들어, X → Y의 통로는 있지만 Y → X의 통로가 없다면 Y는 X로 메시지를 보낼 수 없다. 또한 통로를 통해 메시지를 보낼 때는 일정 시간이 소요된다.

각 도시의 번호와 통로가 설치되어 있는 정보가 주어질 때, 도시 C에서 보낸 메시지를 받게되는 도시의 개수와 도시들이 모두 메시지를 받는 데까지 걸리는 시간은 얼마인지를 계산하는 프로그램을 작성하시오.

✔️ 입력 조건

• 첫째 줄에 도시의 개수 N, 통로의 개수 M, 메시지를 보내고자 하는 도시 C가 공백으로 구분되어 차례대로 주어진다. (1 $\leq$ N $\leq$ 30,000, 1 $\leq$ M $\leq$ 200,000, 1 $\leq$ C $\leq$ N)이 주어진다.
• 둘째 줄부터 M + 1번째 줄에 걸쳐서 통로에 대한 정보 X, Y, Z가 주어진다. 이는 특정 도시 X에서 다른 특정 도시 Y로 이어지는 통로가 있으며, 메시지가 전달되는 시간이 Z라는 의미다. (1 $\leq$ X, Y $\leq$ N, 1 $\leq$ Z $\leq$ 1,000)

✔️ 출력 조건

• 첫째 줄에 도시 C에서 보낸 메시지를 받는 도시의 총 개수와 총 걸리는 시간을 공백으로 구분하여 출력한다.

✔️ 입력 예시

3 2 1
1 2 4
1 3 2

✔️ 출력 예시

2 4

🗒️ 문제 해설

한 도시에서 다른 도시까지의 최단 거리 문제로 치환할 수 있으므로 다익스트라 알고리즘을 이용해서 풀 수 있다.

import sys
import heapq
input = sys.stdin.readline
INF = int(1e9)

n, m, start = map(int, input().split())
graph = [[] for _ in range(n + 1)]
distance = [INF] * (n + 1) # 최단 거리 테이블을 모두 무한으로 초기화

for _ in range(m):
    x, y, z = map(int, input().split())
    graph[x].append((y, z))

q = []
heapq.heappush(q, (0, start))
distance[start] = 0
while q:
    dist, now = heapq.heappop(q)
    if distance[now] < dist:
        continue
    for i in graph[now]:
        cost = dist + i[1]
        if cost < distance[i[0]]:
            distance[i[0]] = cost
            heapq.heappush(q, (cost, i[0]))

count = 0
max_distance = 0
for d in distance:
    if d != INF:
        count += 1
        max_distance = max(max_distance, d)
# 시작 노드는 제외해야 하므로 count - 1
print(count - 1, max_distance)