파이썬으로 다익스트라 알고리즘 구현하기

 

위와 같은 그래프를 노드 1부터 시작해서 최단 거리를 구하는 다익스트라 알고리즘을 구현하면 다음과 같다.

import heapq

def dijkstra(arr, start):
    D = [float("inf")] * (n+1)
    D[start] = 0
    queue = [(0, start)]

    while queue:
        # D값이 가장 작은 노드를 고른다
        now_dist, now_node = heapq.heappop(queue)

        if now_dist > D[now_node]:
            continue

        # 현재 선택한 노드와 인접한 노드들을 읽는다
        for node, weight in arr[now_node]:
            cost = now_dist + weight
            if cost < D[node]:
                D[node] = cost
                heapq.heappush(queue, (cost, node))

    return D[1:]

n = 5
graph = [[],
         [(2, 8), (3, 3)],
         [(4, 4), (5, 15)],
         [(4, 13)],
         [(5, 2)],
         []]

print(dijkstra(graph, 1))

 

실행 결과 [0, 8, 3, 12, 14]가 출력된다.