[알고리즘] 프로그래머스 - 배달

프로그래머스 - 배달

내 풀이

from collections import defaultdict
import sys
import heapq


def dijkstra(start, graph, N, K):
    distances = [sys.maxsize] * (N + 1)
    distances[start] = 0
    heap = []
    heapq.heappush(heap, (0, start))

    while heap:
        cost, node = heapq.heappop(heap)
        for adj_node in graph[node]:
            if distances[adj_node[1]] > cost + adj_node[0]:
                distances[adj_node[1]] = cost + adj_node[0]
                heapq.heappush(heap, (cost + adj_node[0], adj_node[1]))

    count = 0
    for i in distances:
        if i <= K:
            count += 1
    return count


def solution(N, road, K):
    graph = defaultdict(list)
    for edge in road:
        graph[edge[0]].append((edge[2], edge[1]))  # (cost, node)
        graph[edge[1]].append((edge[2], edge[0]))

    return dijkstra(1, graph, N, K)

다익스트라를 heap을 이용해서 더 효율적으로 구현해봤다.