DSA Cheapest Flights within K Stops

The problem

Cheapest Flights within K Stops asks for the lowest fare from a start city to a destination city. The route may use at most K stops. A stop is an intermediate city between the start and the destination, so K stops allow at most K plus 1 flights.

Each flight has a price, and a route costs the sum of its flight prices. A route that needs more stops than K is not allowed, even when it is cheaper. A destination with no allowed route answers minus one.

main = [(0, 1, 100), (1, 2, 100), (2, 0, 100), (1, 3, 600), (2, 3, 200)]
cases = [
    (main, 4, 0, 3, 0, -1),
    (main, 4, 0, 3, 1, 700),
    (main, 4, 0, 3, 2, 400),
    ([(0, 3, 250)], 4, 0, 3, 0, 250),
    ([(0, 1, 50)], 3, 0, 2, 3, -1),
]
for flights, n, src, dst, k, expected in cases:
    print("flights", flights, "src", src, "dst", dst, "k", k, "-> cheapest:", expected)

The first three rows share one flight list and differ only in K. Zero stops allows only a direct flight, and none exists. Two stops unlock the route through cities 1 and 2, which is the cheapest route of all.

Intuition: the cheapest route overall can break the stop limit

The first idea ignores stops and finds the cheapest route by price alone. That route is a real route with a real fare. Its flaw is that nobody checked how many flights it spends.

On the main flight list, the cheapest route by price runs through cities 1 and 2. It uses three flights, which means two stops. At K equals 1 that route is simply not on the table.

import heapq

flights = [(0, 1, 100), (1, 2, 100), (2, 0, 100), (1, 3, 600), (2, 3, 200)]
best = {0: (0, 0)}
heap = [(0, 0, 0)]
while heap:
    price, used, city = heapq.heappop(heap)
    if best.get(city) != (price, used):
        continue
    for u, v, cost in flights:
        if u == city and (v not in best or price + cost < best[v][0]):
            best[v] = (price + cost, used + 1)
            heapq.heappush(heap, (price + cost, used + 1, v))
price, used = best[3]
print("cheapest ignoring stops:", price, "using", used, "flights")
print("stops that route needs:", used - 1)
print("allowed flights at k = 1:", 2)

Price alone answers 400 for a question whose limit forbids that route. The correct answer at K equals 1 is 700. Any method for this problem must count flights while it prices routes, not after.

Brute force: list every route the limit allows

The brute force builds routes flight by flight and never lets a route grow past K plus 1 flights. Every route that reaches the destination is priced. The cheapest price among them is the answer.

A route is extended only from its current last city. The counter records how many partial routes were built. Cycles are pointless here, because every price is positive and a cycle only adds cost.

def brute_cheapest(flights, src, dst, k):
    best = None
    built = 0
    stack = [(src, 0, 0)]
    while stack:
        city, cost, used = stack.pop()
        built += 1
        if city == dst and used > 0:
            if best is None or cost < best:
                best = cost
            continue
        if used == k + 1:
            continue
        for u, v, c in flights:
            if u == city:
                stack.append((v, cost + c, used + 1))
    return (best if best is not None else -1), built

flights = [(0, 1, 100), (1, 2, 100), (2, 0, 100), (1, 3, 600), (2, 3, 200)]
for k in [1, 2]:
    answer, built = brute_cheapest(flights, 0, 3, k)
    print("k =", k, "cheapest:", answer, " partial routes built:", built)

Raising K from 1 to 2 grows the route tree and lowers the fare. Each extra allowed flight multiplies the routes that start with a shared prefix. That multiplication is what the layered method avoids.

The optimal solution: relax in layers from a frozen copy

The price table holds the cheapest known fare to each city, and each round relaxes every flight from a frozen copy of that table, so one round adds at most one flight to any route. That one sentence is the whole method.

Relaxing a flight means testing whether its start price plus its fare beats the known price of its end city. After K plus 1 rounds, the table holds the cheapest fare for every route the limit allows. Freezing is not a detail, because live prices would let one round chain several flights and break the limit it exists to enforce.

def cheapest(n, flights, src, dst, k):
    if src == dst:
        return 0
    prices = [None] * n
    prices[src] = 0
    for _ in range(k + 1):
        frozen = prices[:]
        for u, v, cost in flights:
            if frozen[u] is not None:
                cand = frozen[u] + cost
                if prices[v] is None or cand < prices[v]:
                    prices[v] = cand
    return prices[dst] if prices[dst] is not None else -1

main = [(0, 1, 100), (1, 2, 100), (2, 0, 100), (1, 3, 600), (2, 3, 200)]
cases = [
    (main, 4, 0, 3, 0, -1),
    (main, 4, 0, 3, 1, 700),
    (main, 4, 0, 3, 2, 400),
    ([(0, 3, 250)], 4, 0, 3, 0, 250),
    ([(0, 1, 50)], 3, 0, 2, 3, -1),
]
for flights, n, src, dst, k, expected in cases:
    result = cheapest(n, flights, src, dst, k)
    print("k =", k, "-> cheapest:", result, "expected:", expected, "match:", result == expected)

All five cases pass. The fare only falls as K grows, and it never falls below the stop-blind price of 400 from the intuition section. The limit costs money, and the table shows exactly how much at each K.

Walkthrough: three rounds on the main flight list

The walk runs K equals 2, which means three rounds. The state is the price table after each round. A price of None means no allowed route reaches that city yet.

flights = [(0, 1, 100), (1, 2, 100), (2, 0, 100), (1, 3, 600), (2, 3, 200)]
prices = [0, None, None, None]
print("before any round:", prices)
for r in range(1, 4):
    frozen = prices[:]
    for u, v, cost in flights:
        if frozen[u] is not None:
            cand = frozen[u] + cost
            if prices[v] is None or cand < prices[v]:
                prices[v] = cand
    print("after round " + str(r) + ":", prices)
print("cheapest to city 3:", prices[3])

Round 1 prices only the direct neighbour, city 1. Round 2 reaches city 3 at 700 through city 1. Round 3 starts from the frozen price of city 2 and cuts city 3 to 400. Each round adds exactly one flight of reach.

Complexity

The brute force builds every allowed route, so its count grows exponentially as the flight limit rises, as the counts below show. The layered method runs K plus 1 rounds over E flights, so its time is O(K times E).

The layered method stores one price per city plus one frozen copy, so its space is O(n). The brute force stack holds every partial route it has opened and not yet finished, and that frontier grows with the route count.

def count_routes(n):
    flights = []
    for i in range(n):
        if i + 1 < n:
            flights.append((i, i + 1, 5))
        if i + 2 < n:
            flights.append((i, i + 2, 9))
    routes = 0
    stack = [0]
    while stack:
        city = stack.pop()
        if city == n - 1:
            routes += 1
            continue
        for u, v, c in flights:
            if u == city:
                stack.append(v)
    checks = n * len(flights)
    return routes, checks

for n in [6, 8, 10]:
    routes, checks = count_routes(n)
    print("cities", n, "- brute routes to the end:", routes, " layered relaxation tests:", checks)

The test networks give each city a flight to the next city and a flight over it. Route counts grow like the Fibonacci numbers while relaxation tests grow in a straight line. The layered method pays per flight per round, never per route.

Edge cases and pitfalls

The frozen copy is the pitfall that defines this problem. One live pass over the flights in a lucky order chains three flights inside a single round. That pass reports a fare for a route the limit forbids, and the run below shows it happening.

Zero stops means direct flights only. A start city equal to the destination needs no flight at all, so its fare is zero. Both rows run below through the same layered code.

def cheapest(n, flights, src, dst, k):
    if src == dst:
        return 0
    prices = [None] * n
    prices[src] = 0
    for _ in range(k + 1):
        frozen = prices[:]
        for u, v, cost in flights:
            if frozen[u] is not None:
                cand = frozen[u] + cost
                if prices[v] is None or cand < prices[v]:
                    prices[v] = cand
    return prices[dst] if prices[dst] is not None else -1

print("k = 0, no direct flight:", cheapest(4, [(0, 1, 100), (1, 3, 600)], 0, 3, 0))
print("k = 0, direct flight:", cheapest(4, [(0, 3, 250)], 0, 3, 0))
print("start is the destination:", cheapest(3, [(0, 1, 50)], 0, 0, 2))
print("destination unreachable:", cheapest(3, [(0, 1, 50)], 0, 2, 3))

flights = [(0, 1, 100), (1, 2, 100), (2, 3, 200)]
live = [0, None, None, None]
for u, v, cost in flights:
    if live[u] is not None:
        cand = live[u] + cost
        if live[v] is None or cand < live[v]:
            live[v] = cand
print("one live pass:", live)
frozen_prices = [0, None, None, None]
frozen = frozen_prices[:]
for u, v, cost in flights:
    if frozen[u] is not None:
        cand = frozen[u] + cost
        if frozen_prices[v] is None or cand < frozen_prices[v]:
            frozen_prices[v] = cand
print("one frozen round:", frozen_prices)

The live pass prices city 3 at 400 after what should be a one-flight round. The frozen round leaves city 3 unpriced, which is the only correct state after one flight. Every round in the optimal solution freezes for exactly this reason.