프로그래머스 Lv.2 | Python | 배달

1. 문제 정보

프로그래머스

• 난이도: Lv.2

---

2. 문제 설명

N개의 마을로 이루어진 나라가 있습니다. 이 나라의 각 마을에는 1부터 N까지의 번호가 각각 하나씩 부여되어 있습니다. 각 마을은 양방향으로 통행할 수 있는 도로로 연결되어 있는데, 서로 다른 마을 간에 이동할 때는 이 도로를 지나야 합니다.

도로를 지날 때 걸리는 시간은 도로별로 다릅니다. 현재 1번 마을에 있는 음식점에서 각 마을로 음식 배달을 하려고 합니다. 각 마을로부터 음식 주문을 받으려고 하는데, N개의 마을 중에서 K 시간 이하로 배달이 가능한 마을에서만 주문을 받으려고 합니다. 다음은 N = 5, K = 3인 경우의 예시입니다.

위 그림에서 1번 마을에 있는 음식점은 [1, 2, 4, 5] 번 마을까지는 3 이하의 시간에 배달할 수 있습니다. 그러나 3번 마을까지는 3시간 이내로 배달할 수 있는 경로가 없으므로 3번 마을에서는 주문을 받지 않습니다. 따라서 1번 마을에 있는 음식점이 배달 주문을 받을 수 있는 마을은 4개가 됩니다.

마을의 개수 N, 각 마을을 연결하는 도로의 정보 road, 음식 배달이 가능한 시간 K가 매개변수로 주어질 때, 음식 주문을 받을 수 있는 마을의 개수를 return 하도록 solution 함수를 완성해주세요.

제한사항

• 마을의 개수 N은 1 이상 50 이하의 자연수입니다.
• road의 길이(도로 정보의 개수)는 1 이상 2,000 이하입니다.
• road의 각 원소는 마을을 연결하고 있는 각 도로의 정보를 나타냅니다.
• road는 길이가 3인 배열이며, 순서대로 (a, b, c)를 나타냅니다.
  • a, b(1 ≤ a, b ≤ N, a != b)는 도로가 연결하는 두 마을의 번호이며, c(1 ≤ c ≤ 10,000, c는 자연수)는 도로를 지나는데 걸리는 시간입니다.
  • 두 마을 a, b를 연결하는 도로는 여러 개가 있을 수 있습니다.
  • 한 도로의 정보가 여러 번 중복해서 주어지지 않습니다.
• K는 음식 배달이 가능한 시간을 나타내며, 1 이상 500,000 이하입니다.
• 임의의 두 마을간에 항상 이동 가능한 경로가 존재합니다.
• 1번 마을에 있는 음식점이 K 이하의 시간에 배달이 가능한 마을의 개수를 return 하면 됩니다.

---

3. 접근 기준

이 문제의 핵심은 1번 마을에서 각 마을까지의 최단 거리를 구하는 것이다. 도로에는 이동 비용이 있고, 모든 비용이 동일하지 않다.
따라서 단순 BFS가 아니라 다익스트라 알고리즘을 사용해야 한다.

정리하면 다음 흐름이다.

• 마을과 도로 정보를 그래프로 만든다.
• 1번 마을에서 시작한다.
• 각 마을까지의 최소 이동 시간을 계산한다.
• 계산된 거리 중 K 이하인 마을만 카운트한다.

---

4. 사용한 패턴 / 알고리즘

다익스트라 알고리즘

• 시작 노드에서 다른 모든 노드까지의 최단 거리 계산
• 간선 비용이 있는 그래프에서 사용
• 우선순위 큐를 사용하면 더 짧은 거리부터 탐색 가능

힙(heapq)

• 현재까지 발견한 거리 중 가장 짧은 거리부터 꺼내기 위해 사용

그래프 (인접 리스트)

• 각 마을과 연결된 마을 정보 저장

---

5. 핵심 아이디어

1) dist 배열

dist[i]

의 의미:

1번 마을에서 i번 마을까지의 현재 최단 거리

최소 비용을 구해야기 때문에 처음에는 모두 매우 큰 값으로 초기화한다.

INF = int(1e9)
dist = [INF] * (N + 1)

그리고 시작점만 0으로 설정한다.

dist[1] = 0

 

2) heapq 사용 이유

다익스트라는 현재까지 가장 짧은 거리 부터 탐색해야 한다. 거리를 기준으로 하기 때문에 거리를 첫번째 값으로 배치한다.

그래서 힙에:

(거리, 노드)

형태로 저장한다.

예:

(3, 5)

 

의미:

5번 마을까지 현재 거리 3

heapq는 튜플 첫 번째 값을 기준으로 정렬하므로 가장 짧은 거리부터 꺼낼 수 있다

---

6. 풀이 코드

import heapq

def solution(N, road, K):

    # 그래프 생성
    # graph[a] = [(연결 마을, 비용), ...]
    graph = [[] for _ in range(N + 1)]

    # 양방향 도로 정보 저장
    for a, b, cost in road:
        graph[a].append((b, cost))
        graph[b].append((a, cost))

    INF = int(1e9)

    # 최단 거리 저장 배열
    dist = [INF] * (N + 1)

    # 시작 마을은 1번
    dist[1] = 0

    # 우선순위 큐 생성
    heap = []

    # (현재 거리, 현재 마을)
    heapq.heappush(heap, (0, 1))

    while heap:

        # 현재 가장 짧은 거리의 마을 꺼내기
        cur_dist, cur_node = heapq.heappop(heap)

        # 이미 더 짧은 경로로 처리된 적이 있으면 무시
        if cur_dist > dist[cur_node]:
            continue

        # 현재 마을과 연결된 다음 마을 확인
        for next_node, cost in graph[cur_node]:

            # 다음 마을까지 거리 계산
            next_dist = cur_dist + cost

            # 더 짧은 거리라면 갱신
            if next_dist < dist[next_node]:

                # 최단 거리 갱신
                dist[next_node] = next_dist

                # 힙에 새로운 거리 정보 추가
                heapq.heappush(heap, (next_dist, next_node))

    # K 이하 거리인 마을 개수 반환
    return sum(d <= K for d in dist)

 

6-1. 예시로 이해하기

문제 예시와 같은 구조가 있다고 하자.

 

 

 

도로 및 그래프:

road = [
    [1, 2, 1],
    [1, 4, 2],
    [2, 3, 3],
    [2, 5, 2],
    [4, 5, 2],
    [5, 3, 1]
]

## graph[i]: i번째 마을의 (인접 마을, 비용)
graph = [[], [(2, 1), (4, 2)], [(1, 1), (3, 3), (5, 2)], 
	[(2, 3), (5, 1)], [(1, 2), (5, 2)],
	[(2, 2), (3, 1), (4, 2)]]

 

초기 상태:

dist = [INF, 0, INF, INF, INF, INF]
heap = [(0, 1)]

 

1) 1번 마을 탐색 (탐색 시작 마을)

heap = [(0,1)]
heapq.heappop(heap) # (0, 1) => (비용, 마을)

연결 정보:

graph[1] = [(2, 1), (4, 2)] #(1번 마을의 인접 마을, 비용)

1 → 2 : 비용 1	# dist[2]=INF 이므로 1로 갱신
1 → 4 : 비용 2	# dist[4]=INF 이므로 2로 갱신

갱신 후:

dist = [INF, 0, INF, INF, INF, INF] => dist = [INF, 0, 1, INF, 2, INF]
heap = [(0, 1)] => heap = [(1, 2), (2, 4)]	#(1번 마을로 부터의 거리, 현재 마을)

1번 마을 ->  2번 마을까지 거리 : 1

1번 마을 ->  4번 마을까지 거리 : 2 

현재 가장 짧은 거리: 힙에서 첫번째 값을 기준

heap = [(1,2),(2,4)]
heapq.heappop(heap) #(1, 2)

가 나오므로 2번 마을 탐색

2) 2번 마을 탐색

2번과 연결된 마을 확인: 

graph[2] = [(1, 1), (3, 3), (5, 2)]	#(2번 마을의 인접 마을, 비용)
2 → 1 : 1 + 1 = 2 	# dist[1]=0 이므로 갱신 x
2 → 3 : 1 + 3 = 4	# dist[3]=INF이므로 4로 갱신
2 → 5 : 1 + 2 = 3	# dist[5]=INF이므로 3로 갱신

갱신 후:

dist = [INF, 0, 1, INF, 2, INF] => dist = [INF, 0, 1, 4, 2, 3]
heap = [(2, 4)] => heap = [(2, 4), (4, 3), (3, 5)]

 

즉 다음 탐색은:

heap = [(2, 4), (4, 3), (3, 5)]
heapq.heappop(heap) # (2, 4) => (비용, 마을)

 

가 나오므로 4번 마을 탐색

3) 4번 마을 탐색

4번과 연결된 마을 확인: 

graph[4] = [(1, 2), (5, 2)]
4 → 1 : 2 + 2 = 4	# dist[1]=0 이므로 갱신 X
4 → 5 : 2 + 2 = 4	# dist[5]=3 이므로 갱신 X

 

근데 이미:

dist[1] = 0
dist[5] = 3

 

더 짧은 경로가 있으므로 갱신하지 않는다.

if next_dist < dist[next_node]:

 

조건을 만족하지 못하기 때문에 갱신 안 함, 힙에도 추가 안 함

dist = [INF, 0, 1, 4, 2, 3]

현재 가장 짧은 거리: 힙에서 첫번째 값을 기준

heap = [(3, 5),(4, 3)]
heapq.heappop(heap) #(3,5)

가 나오므로 5번 마을 탐색

 

4) 5번 마을 탐색

5번과 연결된 마을 확인: 

graph[5] = [(2, 2), (4, 2), (3, 1)]

5 → 2 : 3 + 2 = 5   # dist[2]=1 이므로 갱신 X
5 → 4 : 3 + 2 = 5   # dist[4]=2 이므로 갱신 X
5 → 3 : 3 + 1 = 4   # dist[3]=4 이므로 갱신 X

조건을 만족하지 못하기 때문에 갱신 안 함, 힙에도 추가 안 함

dist = [INF, 0, 1, 4, 2, 3]

현재 가장 짧은 거리: 힙에서 첫번째 값을 기준

heap = [(4, 3)]
heapq.heappop(heap) # (4, 3)

가 나오므로 3번 마을 탐색

 

5) 3번 마을 탐색

3번과 연결된 마을 확인: 

graph[3] = [(2, 3), (5, 1)]
3 → 2 : 4 + 3 = 7   # dist[2]=1 이므로 갱신 X
3 → 5 : 4 + 1 = 5   # dist[5]=3 이므로 갱신 X

조건을 만족하지 못하기 때문에 갱신 안 함, 힙에도 추가 안 함

 

최종 결과:

dist = [INF, 0, 1, 4, 2, 3]

heap = []

 

힙이 비었으므로 탐색 종료.

 

K = 3 이라면:

1, 2, 4, 5

총 4개 마을에 배달 가능하다.

---

7. 핵심 로직 요약

이 문제의 핵심은 가장 짧은 거리부터 탐색하면서 최단 거리 배열을 갱신하는 것 이다.

• 그래프를 인접 리스트로 만든다.
• dist 배열에 최단 거리를 저장한다.
• 힙에서 가장 짧은 거리부터 꺼낸다.
• 연결된 마을 거리 갱신
• 더 짧은 경로면 힙에 다시 추가
• 최종적으로 K 이하인 마을 개수 계산

결국, 그래프 + heapq + 다익스트라 구조가 핵심이다.

---

8. 정리

처음에는 BFS처럼 보였는데, 실제로는 비용이 있는 최단 거리 문제 라는 걸 이해해야 했다.

특히 중요한 포인트는 이거였다.

탐색 순서가 노드 번호가 아니라 거리 기준이다

 

그래서 heapq를 사용하는 순간 다익스트라 구조가 만들어진다.

또 하나 느낀 점은 같은 마을이 힙에 여러 번 들어갈 수 있다 는 점이다.

그래서

if cur_dist > dist[cur_node]:
    continue

 

조건이 꼭 필요하다.

결국 이 문제는 다익스트라를 처음 익힐 때

• 왜 heap을 쓰는지
• 왜 거리 배열이 필요한지
• 왜 더 짧은 거리만 갱신하는지

를 이해하기 좋은 대표 문제였다.