최대 1 분 소요

문제 링크

풀이

import heapq

def solution(N, road, K):
    graph = [[] for _ in range(N+1)]
    distance = [float('inf')] * (N+1)
    distance[1] = 0
    
    for a, b, cost in road:
        graph[a].append((b,cost))
        graph[b].append((a,cost))
    
    
    heap = []
    heapq.heappush(heap, (0,1))
    while heap:
        dist , node = heapq.heappop(heap)
        for adj, weight in graph[node]:
            cost = dist + weight
            if cost < distance[adj]:
                distance[adj] = cost
                heapq.heappush(heap, (cost, adj))
                
    answer = 0
    for key in distance:
        if key <= K:
            answer += 1
            
    return answer  
    

풀이 해석

  • 다익스트라 알고리즘을 이용해서 거리를 갱신해나가는 코드
  • heapq를 이용해서 while 문 돌 때마다 최소비용을 뽑아서 전개한다
  • 다익스트라 알고리즘 구현할 때 heapq를 쓰는이유!
  • heapq에서 튜플을 기준으로 했을 때 맨 앞의 값을 기준으로 정렬한다

댓글남기기