-
Notifications
You must be signed in to change notification settings - Fork 14
/
787-cheapest-flights-within-k-stops.py
52 lines (47 loc) · 1.9 KB
/
787-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
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
class Solution:
def findCheapestPrice(self, n: int, flights: List[List[int]], src: int, dst: int, k: int) -> int:
node_dist = [float('inf') for _ in range(n)]
node_dist[src] = 0
for i in range(k + 1):
temp_node_dist = node_dist[:]
for p, q, w in flights:
if node_dist[p] != float('inf') and node_dist[p] + w < temp_node_dist[q]:
temp_node_dist[q] = node_dist[p] + w
node_dist = temp_node_dist
return node_dist[dst] if node_dist[dst] != float('inf') else - 1
# time O(kE), at most O(VE) for original bellman ford algo
# space O(V)
# using graph and bellman ford and shortest path
from heapq import *
class Solution:
def findCheapestPrice(self, n: int, flights: List[List[int]], src: int, dst: int, k: int) -> int:
graph = [[] for _ in range(n)]
for p, q, w in flights:
graph[p].append((q, w))
step_total = k + 1
step_node_dist = [[float('inf') for _ in range(n)] for _ in range(step_total + 1)]
step_node_dist[0][src] = 0
heap = [(0, 0, src)]
heapify(heap)
visited = set()
while heap:
prev, step, node = heappop(heap)
if node == dst:
break
if step == step_total:
continue
if (step, node) in visited:
continue
visited.add((step, node))
for neighbor, w in graph[node]:
d = prev + w
if d < step_node_dist[step + 1][neighbor]:
heappush(heap, (d, step + 1, neighbor))
step_node_dist[step + 1][neighbor] = d
res = float('inf')
for s in range(step_total + 1):
res = min(res, step_node_dist[s][dst])
return res if res != float('inf') else - 1
# time O(Eklog(Ek))
# space O(Vk)
# using graph and dijkstra and shortest path