diff options
Diffstat (limited to 'cheapest_flight_k_stop.py')
-rw-r--r-- | cheapest_flight_k_stop.py | 58 |
1 files changed, 58 insertions, 0 deletions
diff --git a/cheapest_flight_k_stop.py b/cheapest_flight_k_stop.py new file mode 100644 index 0000000..ede631b --- /dev/null +++ b/cheapest_flight_k_stop.py @@ -0,0 +1,58 @@ +# 787. Cheapest Flights Within K Stops +import collections +import heapq + +""" +there are 'n' cities connected by some number of flights. you are given array +'flights' where 'flight[i] = [fromi, toi, pricei]' indicates that there is a +flight from city 'fromi' to city 'toi' with cost 'pricei'. +you are also given three integers 'src, dst, & k'. return the cheapest price +from 'src' to 'dst' with at most 'k' stops. if there is no such route, -1 +""" + + +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 a, b, p in flights: + f[a][b] = p + heap = [(0, src, k + 1)] + while heap: + p, i, k = heapq.heappop(heap) + if i == dst: + return p + if k > 0: + for j in f[i]: + heapq.heappush(heap, (p + f[i][j], j, k - 1)) + return -1 + + +if __name__ == "__main__": + obj = Solution() + print( + obj.findCheapestPrice( + n=4, + flights=[[0, 1, 100], [1, 2, 100], [2, 0, 100], [1, 3, 600], [2, 3, 200]], + src=0, + dst=3, + k=1, + ) + ) + print( + obj.findCheapestPrice( + n=3, flights=[[0, 1, 100], [1, 2, 100], [0, 2, 500]], src=0, dst=2, k=1 + ) + ) + print( + obj.findCheapestPrice( + n=3, flights=[[0, 1, 100], [1, 2, 100], [0, 2, 500]], src=0, dst=2, k=0 + ) + ) |