[Python] 프로그래머스(Lv2) - 배달

안녕하세요 !

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

프로그래머스 배달 문제입니다.

풀이

그래프를 보면, 가중치가 있는 최단경로를 구하는 문제로 볼 수 있습니다.
가중치가 있는 최단경로는 다익스트라 알고리즘을 이용하여 구할 수 있습니다.
이 알고리즘을 이용하여 1번 마을로부터의 거리를 각 마을마다 구해주고, 거리가 K이하인 경우의 수를 구해줍니다.

import heapq


def dijkstra(graph, start):
    distances = {node: float('inf') for node in graph}
    distances[start] = 0
    q = []
    heapq.heappush(q, [distances[start], start])
    
    while q:
        current_distance, current_node = heapq.heappop(q)
        if distances[current_node] < current_distance:
            continue
        
        for adjacent, weight in graph[current_node].items():
            distance = current_distance + weight
            
            if distance < distances[adjacent]:
                distances[adjacent] = distance
                heapq.heappush(q, [distance, adjacent])
    
    return distances
    

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

    road_graph = {}
    
    for r in road:
        a, b, dist = r
        if a not in road_graph:
            road_graph[a] = {}
        
        if b not in road_graph:
             road_graph[b] = {}
                
        if (b in road_graph[a] and dist > road_graph[a][b]):
            continue
            
        road_graph[a][b] = dist
        road_graph[b][a] = dist
    
    distances = dijkstra(road_graph, 1)
    
    for distance in distances.items():
        # print(distance)
        if distance[1] <= K:
            answer += 1
            
    return answer