forked from kamyu104/LeetCode-Solutions
-
Notifications
You must be signed in to change notification settings - Fork 0
/
cheapest-flights-within-k-stops.py
35 lines (31 loc) · 1.09 KB
/
cheapest-flights-within-k-stops.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
# Time: O((|E| + |V|) * log|V|) = O(|E| * log|V|),
# if we can further to use Fibonacci heap, it would be O(|E| + |V| * log|V|)
# Space: O(|E| + |V|) = O(|E|)
import collections
import heapq
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
"""
adj = collections.defaultdict(list)
for u, v, w in flights:
adj[u].append((v, w))
best = collections.defaultdict(lambda: collections.defaultdict(lambda: float("inf")))
min_heap = [(0, src, K+1)]
while min_heap:
result, u, k = heapq.heappop(min_heap)
if k < 0 or best[u][k] < result:
continue
if u == dst:
return result
for v, w in adj[u]:
if result+w < best[v][k-1]:
best[v][k-1] = result+w
heapq.heappush(min_heap, (result+w, v, k-1))
return -1