알고리즘창고

[코드조각] 다익스트라

야호호코코 2021. 5. 7. 17:26
반응형

최소힙을 이용한 다익스트라. 기본으로는 start부터 모든 지점까지의 최소거리가 담긴 list를 반환한다.

import heapq
import math

def dijkstra(nodes, start):
    N = len(nodes)
    
    d = [math.inf] * (N+1)

    q = []
    qlen = 1
    d[start] = 0
    heapq.heappush(q, (d[start], start))

    while qlen > 0:
        tmp = heapq.heappop(q)
        qlen -= 1
        
        if d[tmp[1]] < tmp[0]:
            continue

        for i in nodes[tmp[1]]:
            dtmp = d[tmp[1]] + i[1]
            if d[i[0]] > dtmp:
                d[i[0]] = dtmp
                heapq.heappush(q, (dtmp, i[0]))
                qlen += 1

    return d
반응형