[Level 2] 배달

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

코딩테스트 연습 - 배달

---

접근

1번 마을에서 다른 마을로 배달을 갈 때 K시간 이하가 걸리는 마을의 개수를 구해야 한다. 따라서 하나의 노드에서 다른 모든 노드까지의 최단 거리를 구하는 다익스트라 알고리즘을 이용하여 풀었다. 먼저 마을과 마을 간의 연결 정보가 담긴 road를 통해 그래프를 그려주었다. 

이후 다익스트라 알고리즘을 통해 1번 마을에서 시작해서 다른 마을까지 탐색하여 최단 거리를 구하였다. 그리고 메모제이션을 한 리스트에서 K이하의 시간을 카운트하여 출력한다.

 

구현

- graph에 각각의 마을 정보에 대해 양방향 그래프로 표현하였다.

- dp 리스트는 다른 마을을 방문할 때 최단 거리를 저장해둔다.

    answer = 0
    INF = 10**6
    graph = [[] for _ in range(N+1)]
    dp = [INF] * (N+1)
    for a, b, c in road:
        graph[a].append((b, c))
        graph[b].append((a, c))

- 이제 다익스트라 함수에 출발 마을과 그래프 정보를 전달한다.

- 현재 노드를 기준으로 다음 노드까지의 비용을 계산하고 dp에 저장된 값과 비교하여 새롭게 계산된 값인 next_weight가 작다면 값을 업데이트한다.

- 이 과정을 통해 dp에는 1번 마을에서 다른 마을까지의 최단 거리가 저장되게 된다.

def dijkstra(start, graph):
        nonlocal dp
        hq = []
        heapq.heappush(hq, (start, 0))
        dp[start] = 0
        while hq:
            cur_node, cur_weight = heapq.heappop(hq)
            if dp[cur_node] < cur_weight:
                continue
            
            for next_node, weight in graph[cur_node]:
                next_weight = weight + cur_weight
                if dp[next_node] > next_weight:
                    dp[next_node] = next_weight
                    heapq.heappush(hq, (next_node, next_weight))

- 다익스트라 이후에 dp에 저장된 값 중에 K이하인 값을 카운트한다.

dijkstra(1, graph)
for dist in dp:
	if dist <= K:
		answer += 1

---

전체 코드

import heapq
def solution(N, road, K):
    answer = 0
    INF = 10**6
    graph = [[] for _ in range(N+1)]
    dp = [INF] * (N+1)
    for a, b, c in road:
        graph[a].append((b, c))
        graph[b].append((a, c))
    
    def dijkstra(start, graph):
        nonlocal dp
        hq = []
        heapq.heappush(hq, (start, 0))
        dp[start] = 0
        while hq:
            cur_node, cur_weight = heapq.heappop(hq)
            if dp[cur_node] < cur_weight:
                continue
            
            for next_node, weight in graph[cur_node]:
                next_weight = weight + cur_weight
                if dp[next_node] > next_weight:
                    dp[next_node] = next_weight
                    heapq.heappush(hq, (next_node, next_weight))
    dijkstra(1, graph)
    for dist in dp:
        if dist <= K:
            answer += 1
    

    return answer