다익스트라 알고리즘 파이썬으로 구현하기 - 최단경로

https://www.acmicpc.net/problem/1753

1753번: 최단경로

---

다익스트라 알고리즘을 사용해 최단거리를 구하는 문제입니다. 별다른 고민 없이 알고리즘 그대로 파이썬으로 옮겨 적으면 쉽게 통과할 것 같지만, 한 가지 함정이 있습니다. 바로 빠른 IO 처리입니다.

파이썬에서의 빠른 입출력 구현은 input 대신 sys.stdin.readline을 사용하는 것인데, 입출력이 큰 경우 시간 최적화를 위해 습관화하는 것이 좋습니다. 아래 코드는 우선순위큐를 사용한 다익스트라 알고리즘에 빠른 입출력을 적용하여 시간 초과 없이 통과한 코드입니다. (알고리즘 상에는 변화가 없습니다)

import sys
import heapq

input = sys.stdin.readline    # 빠른 입출력
INF = sys.maxsize

# 다익스트라 알고리즘
def solve(adjacent, K):
    prev = [-1] * (len(adjacent) + 1)    # predecessor
    dist = [INF] * (len(adjacent) + 1)   # source로부터의 최소 거리 배열
    dist[K] = 0

    priority_queue = []
    heapq.heappush(priority_queue, [0, K])

    while priority_queue:
        # 거리가 제일 작은 노드 선택
        current_dist, here = heapq.heappop(priority_queue)

        # 인접 노드 iteration
        for there, length in adjacent[here].items():
            next_dist = dist[here] + length

            if next_dist < dist[there]:
                dist[there] = next_dist
                prev[there] = here
                heapq.heappush(priority_queue, [next_dist, there])

    return dist, prev

if __name__ == "__main__":
    V, E = map(int, input().split())
    K = int(input())
    adjacent = [{} for _ in range(V + 1)]

    for _ in range(E):
        u, v, w = map(int, input().split())

        # 만약 동일한 경로의 간선이 주어질 경우 적은 비용의 간선 선택
        if v in adjacent[u]:
            adjacent[u][v] = min(adjacent[u][v], w)
        else:
            adjacent[u][v] = w

    dist, prev = solve(adjacent, K)

    for d in dist[1:]:
        print(d if d != INF else "INF")