[파이썬/프로그래머스]배달/다익스트라/우선순위큐/가지치기!!!!!!!!!

링크 programmers.co.kr/learn/courses/30/lessons/12978?language=python3

코딩테스트 연습 - 배달

문제 설명

N개의 마을로 이루어진 나라가 있습니다. 이 나라의 각 마을에는 1부터 N까지의 번호가 각각 하나씩 부여되어 있습니다. 각 마을은 양방향으로 통행할 수 있는 도로로 연결되어 있는데, 서로 다른 마을 간에 이동할 때는 이 도로를 지나야 합니다. 도로를 지날 때 걸리는 시간은 도로별로 다릅니다. 현재 1번 마을에 있는 음식점에서 각 마을로 음식 배달을 하려고 합니다. 각 마을로부터 음식 주문을 받으려고 하는데, N개의 마을 중에서 K 시간 이하로 배달이 가능한 마을에서만 주문을 받으려고 합니다. 

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

 

코드

from collections import defaultdict
import heapq

def solution(N, road, K):
    answer = 0
    
    graph = defaultdict(list)
    for u, v, w in road:
        graph[u].append((v, w))
        graph[v].append((u, w))
    
    Q = [(0, 1)]
    dist = defaultdict(int)
    while Q:
        time, node = heapq.heappop(Q)
        if node not in dist:
            dist[node] = time
            for v, w in graph[node]:
                alt = time + w
                heapq.heappush(Q, (alt, v))
    
    for key in dist:
        if dist[key] <= K:
            answer += 1
    
    return answer

노드 간 이동을 할 때 최단 거리로 이동을 해야하기 때문에 다익스트라 알고리즘을 사용해서 풀었다.

dist 딕셔너리에 각 노드로 이동하는데 걸리는 시간을 저장해주고, for-loop를 돌며 K보다 작은 경우에는 방문할 수 있는 마을의 개수 answer에 1을 더해주었다.

 

+ 2021.04.07

네트워크 딜레이 타임 문제를 풀다가 배달 문제에서 프루닝을 안했던게 생각나서 프루닝을 하러왔다.

프루닝을 안했을 경우 가장 오래 걸리는 시간은 테스트 28의 6.99ms다

from collections import defaultdict
import heapq

def solution(N, road, K):
    answer = 0
    
    graph = defaultdict(list)
    for u, v, w in road:
        graph[u].append((v, w))
        graph[v].append((u, w))
    
    Q = [(0, 1)]
    dist = defaultdict(int)
    while Q:
        time, node = heapq.heappop(Q)
        if node not in dist:
            dist[node] = time
            for v, w in graph[node]:
                if v not in dist:
                    alt = time + w
                    heapq.heappush(Q, (alt, v))
    
    for key in dist:
        if dist[key] <= K:
            answer += 1
    
    return answer

프루닝을 해줬다. 자 이제 어떻게 변하나보자.

와 절반이상 시간이 절약됐다. 6.99ms의 28번은 3.13ms가 됐다!!