[프로그래머스 코딩테스트]Level03 - 배달

Level 03 - 배달

문제 설명

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 하면 됩니다.
  

입출력 예

|—|—|—|—| |N |road |K |result| |5 |[[1,2,1],[2,3,3],[5,2,2],[1,4,2],[5,3,1],[5,4,2]] |3 |4| |6 |[[1,2,1],[1,3,2],[2,3,2],[3,4,3],[3,5,2],[3,5,3],[5,6,1]] |4 |4|

풀이

다익스트라 알고리즘을 이용해 풀었다. 간선의 방향이 양방향임을 주의해서 풀어야 한다.
단방향일 때와 다른 부분은

graph[x].append([y, w]) # 노드, 가중치
graph[y].append([x, w])

이부분밖에 없긴 한다.

1. 시작 노드(여기서는 무조건 1)에서 가장 가까운 노드 선택
2. 1에서 선택한 노드에서 갈 수 있는 노드 조사
3. 갈 수 있는 노드들을 조사하여 가중치 계산, 비교
4. 가중치가 현재 가중치 테이블 값보다 작으면 최소 힙과 가중치 테이블 업데이트
5. 최소 힙이 비워질 때까지 반복
   

# 간선의 방향이 "양방향"임을 주의!!!
import heapq
def solution(N, road, K):
    answer = 0
    inf = 1000000000 # 무한대 수
    graph = [[] for _ in range(N+1)] # 노드별로 [현재위치, 가중치]쌍을 저장할 그래프 생성
    heap = [] # graph의 요소들이 담길 heap
    dp = [inf for _ in range(N+1)] # 가중치 테이블
    
    # 초기화(graph 채우기)
    for i in range(len(road)):
        x = road[i][0] # 출발지점
        y = road[i][1] # 도착지점
        w = road[i][2] # 가중치
        graph[x].append([y, w]) # 노드, 가중치
        graph[y].append([x, w])
    
    # 다익스트라 알고리즘
    dp[1] = 0
    heapq.heappush(heap, (1, 0)) # 힙에 시작 노드와 가중치 0(처음에는 무조건 0이므로) push
    while heap: # 힙이 비워져 있지 않다면
        now, weight = heapq.heappop(heap) # 힙에 저장된 현재 노드 위치와 가중치 pop
        for next_node, w in graph[now]: # 그래프의 현재 노드 위치에서 이동 가능한 노드와 해당 가중치
            next_weight = w + weight # 현재 노드에서 다음 노드까지의 가중치 + 현재 노드까지의 가중치
            if next_weight < dp[next_node]: # 현재 노드에서 다음 노드까지의 가중치 + 현재 노드까지의 가중치가 다음 노드의 가중치(초기에 inf값으로 채워진)보다 작으면
                dp[next_node] = next_weight # 가중치 테이블 업데이트
                heapq.heappush(heap, (next_node, next_weight)) # pop했던 heap에 다음 노드와 다음 노드까지의 가중치를 push

    for num in dp: 
        if num <= K: # K보다 작은 가중치를 가지면 카운트 업
            answer+=1
    return answer

print(solution(5, [[1,2,1],[2,3,3],[5,2,2],[1,4,2],[5,3,1],[5,4,2]], 3))