[dijkstra] 백준 1753 파이썬

 

import sys
import heapq

input = sys.stdin.readline
INF = int(1e9)
N, M = map(int, input().strip().split())
start = int(input())


graph = [[] for _ in range(N+1)]
distance = [INF]*(N+1)

for i in range(N):
    a, b, c = map(int, input().strip().split())
    graph[a].append((b,c))
    
def dijkstra(start):
    q = []
    heapq.heappush(q, (0, start))
    distance[start]=0
    
    while q:
    	dist, now = heapq.heappop(q)
        
        if distance[now] < dist:
        	continue
        
        for i in graph[now]:
            cost = dist + i[1]
            if cost < distance[i[0]]:
            	distance[i[0]] = cost
                heapq.heappush(q, (cost, i[0]))

dijkstra(start)

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

graph 배열에는 아래 내용이 포함된다.

A노드에서 B노드로 가는 비용 C

예시) 1번 -> 2번 / 비용 3

        1번 -> 3번 / 비용 1

        2번 -> 4번 / 비용 4

 

distance 배열은 start노드를 기준으로 각 노드까지의 거리이다.

 

처음 heapq에 start노드까지의 거리가 0이라고 넣어준다.

while문은 q배열이 있는 동안 무한 루프가 되며, dist와 now에 heappop을 이용하여 값을 넣어준다.

현재 노드 distance[now]가 꺼낸 dist보다 작다면 해당 노드는 이미 처리가 된 것이다.

 

만약 처리가 되지 않았다면 위에서 저장한 graph배열에서 현재 노드로부터 다른 노드까지의 값들을 가지고 계산을 한다.

예를 들어, graph[1]이라는 배열 안에는 아래의 값들이 포함되어 있다.

(2, 3) -> 2번 노드로 가는데 비용 3

(3, 1) -> 3번 노드로 가는데 비용 1

(4, 4) -> 4번 노드로 가는데 비용 4

for문을 통해 각 요소에 접근하고 dist와 i [1]을 합쳐 비용을 계산한다.

만약 현재 distance[i[0]]보다 비용이 작다면 갱신을 하고 heapq에 넣어준다.