[leetcode] 787. Cheapest Flights Within K Stops @ python

原题

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

解法

参考: [Java/Python] Priority Queue Solution
小根堆+BFS. 用字典储存flights中航班的信息, 然后每次从队列中取出第一个值便是符合题意的最小值, 当当前的城市等于目的地时, 返回价格, 否则遍历当前城市的下一站, 我们将价格, 下一个城市, 剩余步数放入队列.

代码

class Solution(object):
    def findCheapestPrice(self, n, flights, src, dst, k):
        """
        :type n: int
        :type flights: List[List[int]]
        :type src: int
        :type dst: int
        :type K: int
        :rtype: int
        """
        f = collections.defaultdict(dict)
        for u, v, w in flights:
            f[u][v] = w
        # q saves the price to current city, current city, and the remaing steps
        q = [(0, src, k+1)]
        while q:
            price, curr, step = heapq.heappop(q)
            if curr == dst:
                return price
            if step > 0:
                for next in f[curr]:
                    heapq.heappush(q, (price + f[curr][next], next, step-1))
        return -1