백준 알고리즘 1753. 최단경로

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

1753번: 최단경로

 

import sys
# sys.setrecursionlimit(10**7)
import heapq
sys.stdin = open("sample_input.txt", "r")

# 정점의 개수 v, 간선의 개수 e
v, e = map(int, sys.stdin.readline().split())
# 시작 정점 k
k = int(sys.stdin.readline())


graph = {}
for i in range(1, v+1):
    graph[i] = {}

for _ in range(e):
    start, end, w = map(int, sys.stdin.readline().split())
    if end in graph[start]:
        graph[start][end] = min(graph[start][end], w)
    else:
        graph[start][end] = w

# 모든 정점까지의 소요시간을 무한대로 초기화
weights = {vertex: float('infinity') for vertex in graph}
heap = []


def dijkstra_heap(graph, start):
    # 시작 정점의 가중치는 0으로 초기화
    weights[start] = 0
    # (가중치, 현재 정점)
    heapq.heappush(heap, (0, start))

    while heap:
        w, n = heapq.heappop(heap)

        # 현재 정점과 연결된 다른 정점 확인
        for i, j in graph[n].items():
            nw = w + j
            if nw < weights[i]:
                weights[i] = nw
                heapq.heappush(heap, (nw, i))
    return weights
    
weights = dijkstra_heap(graph, k)
for w in weights.values():
    if w == float('infinity'):
        print("INF")
    else:
        print(w)