Cheapest Flights Within K Stops

Problem

https://leetcode.com/problems/cheapest-flights-within-k-stops/

There are n cities connected by some number of flights. You are given an array flights where flights[i] = [from:sub:`i`, to:sub:`i`, price:sub:`i`] indicates that there is a flight from city from:sub:`i` to city to:sub:`i` with cost price:sub:`i`.

You are also given three integers src, dst, and k, return the cheapest price from src to dst with at most k stops. If there is no such route, return -1.

Example 1:

Input: n = 4, flights = [[0,1,100],[1,2,100],[2,0,100],[1,3,600],[2,3,200]], src = 0, dst = 3, k = 1
Output: 700
Explanation:
The graph is shown above.
The optimal path with at most 1 stop from city 0 to 3 is marked in red and has cost 100 + 600 = 700.
Note that the path through cities [0,1,2,3] is cheaper but is invalid because it uses 2 stops.

Example 2:

Input: n = 3, flights = [[0,1,100],[1,2,100],[0,2,500]], src = 0, dst = 2, k = 1
Output: 200
Explanation:
The graph is shown above.
The optimal path with at most 1 stop from city 0 to 2 is marked in red and has cost 100 + 100 = 200.

Example 3:

Input: n = 3, flights = [[0,1,100],[1,2,100],[0,2,500]], src = 0, dst = 2, k = 0
Output: 500
Explanation:
The graph is shown above.
The optimal path with no stops from city 0 to 2 is marked in red and has cost 500.

Constraints:

2 <= n <= 100
0 <= flights.length <= (n * (n - 1) / 2)
flights[i].length == 3
0 <= from:sub:`i`, to:sub:`i`< n
from:sub:`i`!= to:sub:`i`
1 <= price:sub:`i`<= 10:sup:`4`
There will not be any multiple flights between two cities.
0 <= src, dst, k < n
src != dst

Pattern

Dynamic Programming, Depth-First Search, Breadth-First Search, Graph Theory, Heap (Priority Queue)

Approaches

Modified Dijkstra

Explanation

Dijkstra’s algorithm is used to find the shortest path from a source node to any destination node with non-negative edge weights. The key insight is that whenever we visit a node \(u\), we can update the distances to its neighbors \(v\) if reaching \(v\) through \(u\) is cheaper. How do we know the distance to \(u\)? The source node is naturally 0 distance away from itself. From there, we always visit the nearest unvisited node by popping from a min-heap ordered by cost. This guarentees that when we reach a node, we found the optimal cost. When you pop node \(u\) with cost \(C\), every other path to \(u\) must pass through an unvisited node in the heap with \(cost \ge C\), plus some non-negative edge weight on top. So no future path can beat \(C\).

However we need the cheapest price with at most \(k\) stops. A cheaper path using more stops and a costlier path using fewer stops could both be useful later, so we cannot discard either. To handle this, we expand the distance table to two dimensions: dist[node][stops]. When updating neighboring node distances, we increment the stop counter by 1.

dist[neighbor][stops + 1] = w + price

When we push a node onto the heap, we also record the number of hops. If the hops is \(\ge k\), we skip it. Because we are still retrieving the cheapest node, we can be sure that it has the optimal price, regardless of how many stops the other paths used. However, because paths with hops \(\ge k\), some nodes may have prices which would otherwise be suboptimal.

Code

import heapq
from collections import defaultdict


def findCheapestPrice(
    n: int,
    flights: list[list[int]],
    src: int,
    dst: int,
    k: int,
) -> int:
    """Finds the cheapest flight from ``src`` to ``dst`` that is at most ``k``
    stops.
    """
    graph = defaultdict(list)
    for [u, v, price] in flights:
        graph[u].append((v, price))

    dist = [[float("inf")] * (k + 2) for _ in range(n)]

    min_heap = [(0, src, 0)]
    while min_heap:
        price, u, stops = heapq.heappop(min_heap)

        if dst == u:
            return price

        if stops > k or dist[u][stops] < price:
            continue

        if price > dist[u][stops]:
            continue

        for v, w in graph[u]:
            next_price = w + price
            if dist[v][stops + 1] > next_price:
                dist[v][stops + 1] = next_price
                heapq.heappush(min_heap, (next_price, v, stops + 1))

    return -1

Test

>>> from cheapest_flights_within_k_stops__modified_dijkstra import findCheapestPrice
>>> findCheapestPrice(4, [[0,1,100],[1,2,100],[2,0,100],[1,3,600],[2,3,200]], 0, 3, 1)
700
>>> findCheapestPrice(3, [[0,1,100],[1,2,100],[0,2,500]], 0, 2, 1)
200
>>> findCheapestPrice(3, [[0,1,100],[1,2,100],[0,2,500]], 0, 2, 0)
500

Complexity

\(V\) is the number of cities, \(E\) is the length of flights, and \(k\) is the maximum number of stops.

Measure	Complexity	Notes
Time	\(O(E k \log(Ek))\)	each edge can be considered once per stop count, and each heap push and pop costs the log of the heap size
Auxiliary Space	\(O(k(V + E))\)	the `dist` matrix has \(V \times (k + 2)\) entries and the heap can hold up to \(E \cdot k\) entries

cheapest_flights_within_k_stops__modified_dijkstra.findCheapestPrice(n: int, flights: list[list[int]], src: int, dst: int, k: int) → int: Finds the cheapest flight from src to dst that is at most k stops.