[Python] 백준 / gold / 1753번 (최단경로)

문제 링크 : https://www.acmicpc.net/problem/1753

다익스트라 최단경로 알고리즘을 이용하는 문제이다.

문제 조건에서, 최단거리 테이블을 리스트 자료구조로 작성하면 O(V^2) time이 소요되기 때문에 효율성 테스트를 통과할 수 없다는 것을 알 수 있었다.

따라서 heap 자료구조를 이용해서 O(E x logV) time이 소요되는 알고리즘을 작성해야한다.

정답 코드

import sys
import heapq

INF = 987654321

V, E = map(int, sys.stdin.readline().split())
K = int(sys.stdin.readline())
graph = [[] for _ in range(V + 1)]

for _ in range(E):
    u, v, w = map(int, sys.stdin.readline().split())
    graph[u].append((v, w))

distance = [INF] * (V + 1)


def dijkstra(start):
    distance[start] = 0
    q = []
    heapq.heappush(q, (0, start))

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

        for x in graph[now]:
            cost = dist + x[1]

            if cost < distance[x[0]]:
                distance[x[0]] = cost
                heapq.heappush(q, (cost, x[0]))


dijkstra(K)
for i in range(1, V+1):
    if distance[i] == INF:
        print('INF')
    else:
        print(distance[i])