[백준 BOJ_1753] 최단경로 Python 풀이

출처: 백준 온라인 저지

문제

풀이

최단거리 알고리즘인 다익스트라(dijkstra)를 이용하여 구현해주었습니다.

heapq 모듈을 사용하여 구현해주었고, dist를 math.inf로 초기화해주고, 다익스트라를 통해 K에서 출발하여 도착할 수 있는 최단경로를 update해주었습니다.

다익스트라를 다 돌고나면 dist에는 경로가 있는 한 최단거리로 update가 되니, K에서부터 각 정점까지의 최단거리를 출력해주었습니다. update 되지 않고 inf로 남아 있다면 INF를 출력해주었습니다.

코드

from collections import defaultdict
from heapq import heappop, heappush
from math import inf
import sys
input = sys.stdin.readline
def dijkstra(src, adj, V):
    dist = [inf for _ in range(V+1)]
    hq = [[0, src]]
    dist[src] = 0
    while hq:
        w, v = heappop(hq)
        if w > dist[v]:
            continue
        for nw, nv in adj[v]:
            nw += w
            if dist[nv] > nw:
                dist[nv] = nw
                heappush(hq, (nw, nv))
    return dist
V, E = map(int, input().split())
src = int(input())
adj = defaultdict(list)
for _ in range(E):
    u, v, w = map(int, input().split())
    heappush(adj[u], (w, v))
dist = dijkstra(src, adj, V)
for i in range(1,V+1):
    if dist[i] == inf:
        print("INF")
    else:
        print(dist[i])