최단 경로 구하기 - 다익스트라(Dijkstra) (파이썬)

힙을 이용한 최단 경로 알고리즘 구하기

import heapq as hq
def dijkstra(start, graph, distance):
    q = []
    # 시작 노드로 가기 위한 최단 거리는 0으로 설정
    hq.heappush(q, [0, start])
    distance[start] = 0
    while q:
        #최단 거리의 섬과 거리
        dist, now = hq.heappop(q)
        #현재 섬이 처리 되었다면 스킵
        if distance[now] < dist:
            continue
        #현재 섬과 연결된 다른 섬들을 확인
        for e, c in graph[now]:
            #도착지, 비용
            cost = dist + c
            #  현재 노드를 거쳐 다른 섬으로 이동하는 거리가 더 짧은 경우
            if cost < distance[e]:
                distance[e] = cost
                hq.heappush(q,[cost, e])
    return distance

def solution(n, costs):
    answer = 0
    INF = int(10e9)
    distance = [INF] * (n)
    graph = [[] for i in range(n + 1)]

    for s, e, c in costs:
        graph[s].append([e, c])
    distance = dijkstra(0, graph, distance)
    print(distance)
    answer = sum(distance)
    return answer

solution(4, [[0,2,2], [0,1,1],[1,2,5],[1,3,1],[2,3,8]])

0에서 출발한다면 [0, 1, 2, 2]

1번 노드로 가는 비용 : 1

2번 노드로 가는 비용 : 2

3번 노드로 가는 비용 : 2