[python/leetcode]cheapest flights within k stops/dijkstra

💬 문제

(leetcode.com/problems/cheapest-flights-within-k-stops/)

There are n cities connected by m flights. Each flight starts from city u and arrives at v with a price w.

Now given all the cities and flights, together with starting city src and the destination dst, your task is to find the cheapest price from src to dst with up to k stops. If there is no such route, output -1.

 

💬 코드

class Solution:
    def findCheapestPrice(self, n: int, flights: List[List[int]], src: int, dst: int, K: int) -> int:
        graph = collections.defaultdict(list)
        for s, d, p in flights:
            graph[s].append((d, p))
        
        hq = [(0, src, K)]
        while hq:
            price, node, k = heapq.heappop(hq)
            # 목적지 도착했으면 가격 리턴
            if node == dst:
                return price
            # 경유 횟수(k>=0 이면 아직 경유해도 되는 상태)
            if k >= 0:
                # s(출발지)랑 이어져있는 d(경유지)와 그 곳으로 이동하기 위해 필요한 가격 p
                for d, p in graph[node]:
                    heapq.heappush(hq, (price+p, d, k-1))
        # 경로가 존재하지 않을 경우 -1 리턴
        return -1