aboutsummaryrefslogtreecommitdiff
path: root/cheapest_flight_k_stop.py
diff options
context:
space:
mode:
Diffstat (limited to 'cheapest_flight_k_stop.py')
-rw-r--r--cheapest_flight_k_stop.py58
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
+ )
+ )