[프로그래머스] 배달

문제 링크

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

문제 설명

• 에지 리스트 주어짐
• 1번 노드에서 최단 거리가 K 이하인 노드 개수 구하기

풀이

• 다익스트라
  • 처리안된 노드들 중 가장 가까운 노드 heap에서 pop
  • 인접 노드들 거리 갱신 후 push
  • 반복

후기

• 다익스트라는 봐도봐도 까먹는다

코드

import heapq

def solution(N, road, K):
    answer = 0
    
    # 인접리스트 생성
    edges = [[] for _ in range(N+1)]
    for s,e,d in road:
        edges[s].append([e,d])
        edges[e].append([s,d])

    # 최단거리 저장할 리스트
    dist = [float("inf")] * (N+1)
    dist[1] = 0 
    
    # 처리여부
    checked = [0] * (N+1)
    
    # 현재까지 최단 거리 노드를 위한 힙 
    q = [[0, 1]]
    
    # 다익스트라
    while q:
        cd, s = heapq.heappop(q)                # 방문하지 않은 최단 노드 꺼내기 
        if checked[s]:                          # 방문했으면 continue
            continue
        checked[s] = 1                          # 방문 체크
        for e, nd in edges[s]:                  # 인접 노드 갱신
            dist[e] = min(dist[e], cd + nd)
            heapq.heappush(q, [dist[e], e])     # 일단 넣음. 최단이면 나오겠지~ 아니면 말고
            
    # 개수 세기
    for d in dist:
        if d <= K:
            answer += 1

    return answer