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 airport
  • dst is the destination airport
  • src != dst
  • k is the maximum number of stops you can make (not including src and dst)

Return the cheapest price from src to dst with at most k stops, or return -1 if it is impossible.

Example 1:

# Example 1
n = 4
flights = [[0, 1, 200], [1, 2, 100], [1, 3, 300], [2, 3, 100]]
src = 0
dst = 3
k = 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 2
n = 3
flights = [[1, 0, 100], [1, 2, 200], [0, 2, 100]]
src = 1
dst = 2
k = 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.

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)