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
53
54
55
56
57
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
)
)
|