21. 최단경로문제(2)

https://leetcode.com/problems/cheapest-flights-within-k-stops/

Cheapest Flights Within K Stops - LeetCode

 

 시작점에서 도착점까지의 가장 저렴한 가격을 계산하되, K개의 경유지 이내에 도착하는 가격을 리턴하라. 경로가 존재하지 않을 경우 -1을 리턴한다.

 

 여기서는 앞에 문제와 비교했을때 한 가지 제약사항이 추가되었는데 K개의 경유지 이내에 도착해야 한다는 점이다. 따라서 다익스트라 알고리즘의 구현을 위해 우선순위 큐에 추가할 때 K 이내일 때만 경로를 추가하여 K를 넘어서는 경로는 더 이상 탐색되지 않게 하면 될 것 같다.

 

import collections
import heapq

class Solution:
    def findCheapestPrice(self, n: int, flights:list, src: int, dst: int, K: int):
        graph=collections.defaultdict(list)
        for u,v,w in flights:
            graph[u].append((v,w))    

        k=0
        Q=[(0,src,k)]
        
        while Q:
            price,node,k=heapq.heappop(Q)
            if node==dst:
                return price
            if k<=K:
                k+=1
                for v,w in graph[node]:
                    alt=price+w
                    heapq.heappush(Q,(alt,v,k))
                    
        return -1