프로그래머스 - 배달 (파이썬)

https://school.programmers.co.kr/learn/courses/30/lessons/12978

프로그래머스

 

풀이

N개의 마을(1번부터 N번), roads는 [마을 A, 마을 B, 걸리는 시간] 형태의 리스트고, 음식 배달은 1번 마을에서 시작한다.

K시간 이하로 배달 가능한 마을 수를 구해야한다.

1번 마을을 시작점으로 해서 모든 마을까지의 최단 거리를 구한다. 여기서는 다익스트라 알고리즘을 사용한다.

구한 거리 중에서 K이하인 마을 개수를 세면 정답이다.

시간 복잡도 분석

도로수 E, 마을 수 N

즉, O((N + E) Log N)

코드

import heapq
from collections import defaultdict

def solution(N, road, K):
    G = defaultdict(list)
    for a, b, cost in road :
        G[a].append((b, cost))
        G[b].append((a, cost))
    
    dist = [float('inf')] * (N + 1)
    dist[1] = 0
    
    heap = [(0, 1)]
    while heap :
        curr_dist, curr_node = heapq.heappop(heap)
        
        if curr_dist > dist[curr_node] :
            continue
        
        for neighbor, weight in G[curr_node] :
            distance = curr_dist + weight
            if distance < dist[neighbor] :
                dist[neighbor] = distance
                heapq.heappush(heap, (distance, neighbor))
        
    return sum(1 for d in dist if d <= K)