다익스트라 알고리즘 feat. heap

Yodi.Song·2021년 7월 3일
0

근데이제 프로그래머스 문제 배달을 곁들인

import heapq as h

def dijkstra(graph, start, N, K):
    dstncs = [K+1] * (N+1)
            
    dstncs[start] = 0
    queue = []
    h.heappush(queue, [start, dstncs[start]]) # node, distance
    
    while queue:
        node, dstnc = h.heappop(queue)

        if dstncs[node] < dstnc: continue

        
        for nxt_node, nxt_dstnc in graph[node].items():
            new_dstnc = dstnc + nxt_dstnc
            if new_dstnc < dstncs[nxt_node]:
                dstncs[nxt_node] = new_dstnc
                h.heappush(queue, [nxt_node, new_dstnc])
                
    return dstncs
    

  
    

def solution(N, road, K):
    answer = 0

    # 그래프 만들기
    graph = {i+1 : dict() for i in range (N)}
    
    
    for n1, n2, v in road:
        if n2 in graph[n1]:
            if v > graph[n1][n2]:
                continue
        graph[n1].update({n2 : v})
        graph[n2].update({n1 : v})
        
#     print(graph)
    
    distances = dijkstra(graph, 1, N, K)
    # print(distances)
    
    for node in range(1, N+1):
        if distances[node] <= K:
            answer += 1
    
    return answer

0개의 댓글