[백준] 1753번 최단경로

백준_1753

시작 노드에서, 모든 노드로까지의 최단 거리를 구하는 문제이므로 다익스트라 알고리즘을 사용한다

import heapq
INF = 1e9
v, e = map(int, input().split())
start = int(input())
graph = [[] for _ in range(v+1)]
visited = [0]*(v+1)
distance = [INF]*(v+1)

for _ in range(e):
    a, b, c = map(int, input().split())
    graph[a].append((b,c))  #(끝점, 거리) append

def dijkstra(start):
    q = []
    #(지금까지의 거리, 노드) push
    heapq.heappush(q, (0,start))
    distance[start] = 0

    while q:
        #거리가 짧은 순으로 pop
        dist, now = heapq.heappop(q)
        #distance[now] < dist = 이미 방문되었다는 뜻이므로 무시하고 진행
        if distance[now] < dist:
            continue
        for e,c in graph[now]:
            cost = dist+c
            if cost < distance[e]:
                distance[e] = cost
                heapq.heappush(q, (cost, e))

dijkstra(start)
for i in range(1,v+1):
    if distance[i]==INF:
        print("INF")
    else:
        print(distance[i])

Python3은 시간초과, Pypy3는 성공