Cheapest Flights Within K Stops
Problem Statement
There are n
airports, labeled from 0
to n - 1
, which are connected by some flights. You are given an array flights
where flights[i] = [from_i, to_i, price_i]
represents a one-way flight from airport from_i
to airport to_i
with cost price_i
. You may assume there are no duplicate flights and no flights from an airport to itself.
You are also given three integers src
, dst
, and k
where:
src
is the starting airportdst
is the destination airportsrc != dst
k
is the maximum number of stops you can make (not includingsrc
anddst
)
Return the cheapest price from src
to dst
with at most k
stops, or return -1
if it is impossible.
Example 1:
# Example 1n = 4flights = [[0, 1, 200], [1, 2, 100], [1, 3, 300], [2, 3, 100]]src = 0dst = 3k = 1
# Output: 500
Explanation: The optimal path with at most 1 stop from airport 0 to 3 is shown in red, with total cost 200 + 300 = 500
. Note that the path [0 -> 1 -> 2 -> 3]
costs only 400, and thus is cheaper, but it requires 2 stops, which is more than k.
Example 2:
# Example 2n = 3flights = [[1, 0, 100], [1, 2, 200], [0, 2, 100]]src = 1dst = 2k = 1
# Output: 200
Explanation: The optimal path with at most 1 stop from airport 1 to 2 is shown in red and has cost 200
.
Constraints:
1 <= n <= 100
fromi != toi
1 <= pricei <= 1000
0 <= src, dst, k < n
You should aim for a solution as good or better than O(n + (m * k))
time and O(n)
space, where n
is the number of cities, m
is the number of flights, and k
is the number of stops.
Recommended Time and Space Complexity
You should aim for a solution as good or better than O(n + (m * k))
time and O(n)
space, where n
is the number of cities, m
is the number of flights, and k
is the number of stops.
Hint 1
Consider this as a graph problem where the cities are nodes and the flights are edges connecting two cities, with the ticket cost as the edge weight. Can you think of a shortest path algorithm to solve the problem? Perhaps a better algorithm than Dijkstra’s that can intuitively handle the k
stops condition.
Hint 2
We can use the Bellman-Ford algorithm. Initialize a prices
array of size n
with Infinity
, setting prices[source] = 0
. These values describe the cost to reach a city from the source city. Iterate (k + 1)
times (stops are 0-indexed), updating the cost to each city by extending paths from cities with valid costs. We only update the cost for a city if it is less than the previous cost. How would you implement this?
Hint 3
At each level of iteration, we go through the given flights and use them to update the price array with the minimum costs compared to the previous level. We use a temporary prices array at each level to store the updated costs. After completing all levels, we return the result stored in prices[dst]
. If that value is Infinity
, we return -1
instead.
Solution
Dijkstra’s Algorithm
import heapq
def find_cheapest_price_dijkstra(n: int, flights: list[list[int]], src: int, dst: int, k: int) -> int: inf = float("inf") adj = [[] for _ in range(n)] dist = [[inf] * (k + 5) for _ in range(n)] for u, v, cst in flights: adj[u].append([v, cst])
dist[src][0] = 0 min_heap = [(0, src, -1)] # cost, node, stops while min_heap: cst, node, stops = heapq.heappop(min_heap) if dst == node: return cst if stops == k or dist[node][stops + 1] < cst: continue for nei, w in adj[node]: next_cst = cst + w next_stops = 1 + stops if dist[nei][next_stops + 1] > next_cst: dist[nei][next_stops + 1] = next_cst heapq.heappush(min_heap, (next_cst, nei, next_stops))
return -1
Time complexity: O((n + m) * k)
Space complexity: O(n * k)