[프로그래머스] lv2 배달_python

이번 문제는 각 노드의 관계가 주어지고, 1번 마을로부터 각 마을까지의 최소거리를 구하여 원하는 결과를 도출한다는 점에서 다익스트라 또는 플로이드 워셜을 생각했다.

다익스트라와 힙구조를 사용하였다.

거리를 담을 dist 리스트를 생성해주고, 값은 inf로 초기화 했다.(최소거리를 구하기 때문)
관계를 담을 graph 리스트를 생성하였다. graph에 담기는 순서는 [ 노드, 배달시간 ] 순서이다.

다익스트라 알고리즘을 작성하는데,
heap 을 생성하고 그 안에 시작 값을 담았다. [ 1, 0 ] <- 1번 마을에서 1번마을까지는 0시간.
이후 반복문을 통해 heap이 빌때까지, 각 노드의 관계와 시간을 꺼내면서 시간을 비교하여 작은것을 담는 동시에 경유할때 발생하는 시간 또한 포함하였다. dist에 값을 계속해서 변경하였고, 이후 다시 heap에 담고 이 동작을 반복.

답을 도출할땐, dist안의 값에서 목표 시간 K와 같거나 작은 값만을 카운팅하여 답을 도출.
(편의상 0번 인덱스에는 inf값을 담아 리스트를 마을 수 보다 1개 더 많게 설정)

import sys
import heapq as hq

def dij(dist, graph):
    heap = []
    hq.heappush(heap, [1,0])
    while heap:
        node, cost = hq.heappop(heap)
        for n, c in graph[node]:
            if cost+c < dist[n]:
                dist[n] = cost + c
                hq.heappush(heap, [n, cost+c])
            
def solution(N, road, K):
    answer = 0
    # N 마을 수, road 관계, K 목표시간
    dist = [float('inf') for _ in range(N+1)]
    dist[1] = 0
    graph = [[] for _ in range(N+1)]
    for start, end, cost in road:
        graph[start].append([end, cost])
        graph[end].append([start, cost])
    dij(dist, graph)
    answer = len([i for i in dist if i <= K])
    return answer