Cheapest flights within K stops in Python

Why it matters

A travel site is asked for the cheapest flight from one city to another, with a rule from the customer: at most one stop. The cheapest flight overall might have three stops, and the customer does not want it. The site must find the cheapest route among the routes that respect the limit. This is a shortest path problem with one extra rule. The route is limited by the number of flights it uses, and the cost is the sum of the prices.

Two earlier lessons matter here. The Dijkstra lesson found the cheapest route in a weighted graph, but it cannot respect a limit on the number of edges, because it always keeps only the best price to each city. The Bellman-Ford part of that lesson handled negative weights by running rounds of relaxation. This lesson uses those rounds, but stops after a fixed number of them. Each round allows one more flight, so the number of rounds sets the limit.

This lesson covers one problem with two methods. Both methods produce the same answers, and both need a copy of the distances for each round. The copy is the detail that most beginners get wrong.

The programs in this lesson put the flight map in the code, in a list called FLIGHTS. They do not read input, so each one runs as it is. To try another map, change the list.

Words used in this lesson

A flight is one edge: a direct trip from one city to another, with a price. In this lesson a flight always goes one way, so the graph is directed. In Python a flight is a tuple of three values, (from_city, to_city, price).

A stop is a city where you change planes. A route with one flight has zero stops. A route with two flights has one stop, because the city between them is a stop. In general, a route with f flights has f - 1 stops.

The limit K is the largest number of stops allowed. A route with at most K stops uses at most K + 1 flights. This off-by-one step is the source of most bugs in this problem, so keep the two numbers apart in your head: K counts stops, and the program counts flights.

A relaxation is the check "is there a cheaper way to reach this city now?" If the best known price to city u is d, and the flight from u to v costs w, then the new candidate price for v is d + w. If that is smaller than the best known price to v, update it.

A round is one pass over every flight, relaxing each one once.

Why the cheapest route is not always the simplest

Take four cities, numbered 0 to 3. We want to fly from city 0 to city 3.

• Flight 0 to 1 costs 10.
• Flight 1 to 2 costs 10.
• Flight 2 to 3 costs 10.
• Flight 0 to 2 costs 50.
• Flight 0 to 3 costs 100.

The direct flight from 0 to 3 costs 100 and has zero stops. The route 0, 2, 3 costs 50 plus 10, which is 60, and has one stop. The route 0, 1, 2, 3 costs 10 plus 10 plus 10, which is 30, and has two stops.

If the limit is zero stops, the answer is 100. If the limit is one stop, the answer is 60. If the limit is two stops, the answer is 30. On this map the cheapest route overall uses the most stops, and a customer who allows only one stop gets a different answer.

This is why the limit cannot be handled by changing the prices. We cannot simply pick the cheapest route and then check its stops, because the cheapest route might have too many stops, and a slightly more expensive route with fewer stops might be the right answer. We need to track the number of flights as part of the search.

The idea: rounds that allow one more flight

Think about the question in rounds. After round 1, we know the cheapest price to each city using at most one flight. After round 2, we know the cheapest price using at most two flights. After round r, we know the cheapest price using at most r flights. Each round extends the set of routes by one flight.

Start with dist[src] = 0, and every other city at infinity. In Python the value math.inf means "not reached yet". It is a float larger than every price. Adding a price to math.inf gives math.inf, so an unreached city never looks cheap. The programs still check for it, so the reader can see where unreached cities are skipped.

In round 1, each flight that starts at a city we have reached can improve its destination. The only reached city at the start is the source. So round 1 handles the flights out of the source. In round 2, the cities reached in round 1 can send flights, and so on.

The answer after K + 1 rounds is the cheapest price to the destination, using at most K + 1 flights. That is exactly the limit on stops.

Worked example: the four-city map, round by round

Use the four-city map above with K = 1. That means at most two flights, so we run K + 1 = 2 rounds. The flights are listed in this order: 0 to 1, 1 to 2, 2 to 3, 0 to 3, 0 to 2. Here inf stands for infinity.

Start: the distances are [0, inf, inf, inf].

Round 1 reads each flight with the distances from the start of the round. The only city with a finite distance is 0, so only flights out of 0 can help.

• Flight 0 to 1 costs 10. The candidate price to 1 is 0 + 10 = 10, so city 1 becomes 10.
• Flight 0 to 3 costs 100, so city 3 becomes 100.
• Flight 0 to 2 costs 50, so city 2 becomes 50.
• The flights from 1 and 2 are skipped, because those cities are still at infinity at the start of the round.

After round 1 the distances are [0, 10, 50, 100]. These are the cheapest prices using at most one flight.

Round 2 uses the distances from the end of round 1. Now city 1 has price 10 and city 2 has price 50.

• Flight 0 to 1 gives 0 + 10 = 10, no change.
• Flight 1 to 2 gives 10 + 10 = 20, which is less than 50, so city 2 becomes 20.
• Flight 2 to 3 gives 50 + 10 = 60, which is less than 100, so city 3 becomes 60.
• Flight 0 to 3 gives 0 + 100 = 100, no change.
• Flight 0 to 2 gives 0 + 50 = 50, no change, because city 2 is now 20.

After round 2 the distances are [0, 10, 20, 60]. The answer for city 3 is 60, which matches the one-stop answer from the map.

Notice the detail in round 2. The flight from 1 to 2 used the price of city 1 from the end of round 1. The flight from 2 to 3 used the price of city 2 from the start of round 2, which was 50, not the 20 that appeared earlier in the same round. The rule that makes this work is the copy of the distances. Read the old values, write the new values into a copy, and then replace the old list with the copy.

Why a copy of the distances is needed

Here is the failure. Suppose a program updates the distance list in place, so every write is visible to the next read in the same round. The flights are read in the order given. Flight 0 to 1 sets city 1 to 10. Flight 1 to 2 then reads city 1 as 10 and sets city 2 to 20. Flight 2 to 3 then reads city 2 as 20 and sets city 3 to 30. All three flights were used in one round.

One round is supposed to add one flight. The in-place program used three flights in a single round, so it lets a route pass through as many stops as the flight list happens to allow. The order of the flights in the input decides the answer, and that is a serious bug.

The fix is the copy. Each round starts with nxt = dist.copy(). Every relaxation reads from dist, which holds the values from the end of the previous round, and writes into nxt. At the end of the round, dist = nxt. Now a round can add at most one flight to any route, no matter how the input is ordered.

The programs in this lesson show the copy and the bug. The mistakes section runs the buggy version on the four-city map and shows the wrong answer it gives.

Program 1: Bellman-Ford with K + 1 rounds

This program runs K + 1 rounds, each with a copy of the distances, and returns the answer or -1 if the destination cannot be reached within the limit. The loop for u, v, price in flights unpacks each tuple into three names. That is the Python way to read the three parts of a flight.

The output is:

K = 0: 100
K = 1: 60
K = 2: 30

With K = 1 the program runs two rounds, so it allows two flights. The route 0, 2, 3 uses two flights: 50 from 0 to 2, then 10 from 2 to 3. That is 60, and it is cheaper than the direct flight at 100. The route 0, 1, 2, 3 needs three flights, so it is not allowed.

The dist.copy() call makes the new list nxt. The name nxt is short for "next round". A slice dist[:] makes the same copy, and .copy() says what it does more clearly.

Idea 2: a search by levels

The second method is a breadth-first search by levels. A breadth-first search (BFS) visits the cities in rings: first the cities one flight away from the source, then those two flights away, and so on. Each ring is one level, and one level is one flight.

The search keeps a set called the frontier. The frontier holds the cities whose price improved in the last level. A city that did not improve cannot give a new price, because its old price was already passed on in an earlier level.

At each level, the program does the following. It creates nxt, a copy of the best prices so far. For each city in the frontier, it tries each flight out of that city, using the price from the start of the level. If the new price is lower than the price in nxt, it updates nxt and adds the destination to the new frontier. At the end of the level, best becomes nxt, and the frontier becomes the new set of improved cities.

The stopping rule is simple. The loop ends when the level count reaches K + 1, or when the frontier is empty. An empty frontier means no price improved in the last level, so no later level can improve anything either.

Here is the same map run by levels with K = 1:

• The first level (it uses one flight) starts with the frontier of city 0. The flights from 0 give prices 10 to city 1, 50 to city 2, and 100 to city 3. The new frontier is cities 1, 2 and 3. The prices are [0, 10, 50, 100].
• The second level (it uses two flights) starts with the frontier of cities 1, 2 and 3. Only cities 1 and 2 have flights out. The flight from 1 to 2 gives 20, which improves city 2. The flight from 2 to 3 gives 60, which improves city 3. The prices are [0, 10, 20, 60].

The answer is 60 for one stop. The frontier method gives the same answers as the rounds method, but it does less work, because it skips cities that did not improve. The work is still about K * E steps in the worst case, where E is the number of flights, because each level can scan every flight once.

Program 2: a search by levels

This program builds an adjacency list for each city. An adjacency list is a list, one entry per city, that holds the flights leaving that city. The expression adj[u] gives the flights out of city u. The frontier is a Python set, written with braces, and {src} is a set holding one city.

The output is:

K = 0: 100
K = 1: 60
K = 2: 30

Look at the line candidate = best[u] + price. It reads best, not nxt. This is the same copy rule as before. The level reads the prices from the start of the level and writes the new prices into nxt. Changing it to read nxt would let one level use two flights, which is the bug from the earlier section.

The order in which a set gives its cities does not change the answer, because every read uses best. A program must not depend on that order, even though CPython happens to keep small integers in a predictable order.

The two programs were checked against each other on 60 random small graphs, with 2 to 6 cities, 1 to 10 flights, weights from 1 to 20, and limits from 0 to 4. They gave the same answer in all 60 cases.

A harder example

The four-city map was small. Here is a map with more cities and a case where the best answer comes from a route that uses one extra stop.

Take five cities, numbered 0 to 4. The flights are:

• 0 to 1, price 5
• 1 to 2, price 5
• 2 to 4, price 5
• 0 to 3, price 2
• 3 to 4, price 40
• 0 to 4, price 100

We want to fly from 0 to 4. The cheapest route with no limit is 0, 1, 2, 4, which costs 15 and has two stops. The route 0, 3, 4 costs 42 and has one stop. The direct flight costs 100 and has no stops.

With K = 0, the answer is 100. With K = 1, the answer is 42. With K = 2, the answer is 15. Program 1 gives those three values when its flight list, source and destination are changed to this map. The cheapest route changes as the limit grows, and a program that always takes the globally cheapest path would be wrong for K = 1.

Work through K = 1 with the rounds method by hand. Round 1 reads the flights out of city 0: city 1 gets 5, city 3 gets 2, and city 4 gets 100. The distances are [0, 5, inf, 2, 100]. Round 2 uses those values. The flight from 1 to 2 gives city 2 the price 10. The flight from 3 to 4 gives 2 + 40 = 42, which is less than 100, so city 4 becomes 42. The flight from 2 to 4 would need city 2 to have a price from round 1, and it does not, so it adds nothing in round 2. Final answer: 42.

Now try K = 2. Round 3 uses the distances after round 2. City 2 now has price 10. The flight from 2 to 4 gives 10 + 5 = 15, which is less than 42, so city 4 becomes 15. Final answer: 15. This is the cheapest route that uses at most two stops.

The same five-city map also shows why the copy matters. In round 2 the flight from 2 to 4 must not see the price of city 2 that was written earlier in the same round, because city 2 only became reachable in that round. With the copy, the flight waits for round 3, which is what the limit requires.

Common mistakes

Mistake 1: updating the distances in place. This is the bug the copy prevents. The program below is the rounds method without the copy() call. It reads and writes the same list. Here is the four-city map with K = 0, which should give 100 because only one flight is allowed.

The output is:

30

The answer should be 100, the direct flight, because zero stops allows one flight only. The buggy program prints 30, a route with three flights. In one round it chained flight 0 to 1, flight 1 to 2 and flight 2 to 3, because each flight read the value the previous flight had just written. The program allowed three flights while the limit said one.

Mistake 2: running K rounds instead of K + 1. The loop must run once for each allowed flight. Change the loop in Program 1 from range(k + 1) to range(k). Then run it again. The output is:

K = 0: -1
K = 1: 100
K = 2: 60

The correct answer for K = 1 is 60. The program printed 100 because one round can reach only cities one flight away from the source, so the two-flight route 0, 2, 3 is never tried. With K = 0 it runs no rounds at all, so no city is reached and the answer is -1.

Mistake 3: expecting integers to overflow. In a language with fixed-size integers, a sum can wrap around to a negative number. Python integers have no largest value, so a sum of large prices stays exact. These two lines show it:

The output is:

3000000000
1267650600228229401496703205376

The wrap-around bug does not happen here. The sentinel math.inf is not a whole number, so adding a price to it gives math.inf. Mistake 4 shows what goes wrong when a program uses None for "not reached".

Mistake 4: using None for "not reached". Some programs use None in place of math.inf. Python does not allow arithmetic on None, so the program stops with an error:

dist = [0, None, None]
candidate = dist[1] + 10

The output is:

Traceback (most recent call last):
  File "example.py", line 2, in <module>
    candidate = dist[1] + 10
                ~~~~~~~~^~~~
TypeError: unsupported operand type(s) for +: 'NoneType' and 'int'

The file name, line numbers and caret marks depend on your file and Python version. The fix is to use math.inf, which adds and compares like a number.

Mistake 5: a city number outside the range. The city numbers must be from 0 to n - 1. The program below says there are three cities, but its flight list names city 3, which does not exist:

import math
 
INF = math.inf
FLIGHTS = [(0, 1, 10), (1, 2, 10), (2, 3, 10)]
 
dist = [INF, INF, INF]
dist[0] = 0
for u, v, price in FLIGHTS:
    if dist[u] == INF:
        continue
    dist[v] = min(dist[v], dist[u] + price)
print(dist)

The output is:

Traceback (most recent call last):
  File "example.py", line 11, in <module>
    dist[v] = min(dist[v], dist[u] + price)
                  ~~~~^^^
IndexError: list index out of range

The message says that the program asked for an index that the list does not have. The file name and line numbers depend on your file. The fix is to make the city numbers match n.

Practice

Try each exercise by hand first, then check with the programs.

Exercise 1 (easy). Use the five-city map from the harder example. Write the distance list after round 1 when the source is city 0.

Show the answer

The distances are [0, 5, inf, 2, 100]. Only the flights out of city 0 can be used in round 1. Those are 0 to 1 (price 5), 0 to 3 (price 2) and 0 to 4 (price 100). City 2 is still at infinity.

Exercise 2 (easy). For the four-city map, what is the answer with K = 0? Then what is the answer with K = 2?

Show the answer

With K = 0 the answer is 100, the direct flight from 0 to 3. With K = 2 the answer is 30, the route 0, 1, 2, 3 with three flights and two stops.

Exercise 3 (medium). Three cities, numbered 0 to 2. Flights: 0 to 1 costs 5, 1 to 2 costs 5, and 0 to 2 costs 20. Find the answer for K = 0 and for K = 1.

Show the answer

With K = 0 the answer is 20, the direct flight. With K = 1 the answer is 10, the route 0, 1, 2. Both programs in this lesson give these values for this map.

Exercise 4 (medium). Explain in your own words why the search by levels can stop early when the frontier is empty, and why that stop does not change the answer.

Show the answer

The frontier holds the cities whose price improved in the last level. If it is empty, no price changed in the last level. The next level would read the same prices and make the same changes, which is none. So every later level would also change nothing, and the answer is already final.

Exercise 5 (hard). The rounds method runs K + 1 rounds, and each round scans every flight. A map has n cities and E flights. What is the running time in terms of K and E? Does the time depend on n?

Show the answer

Each round scans E flights once, and there are K + 1 rounds. The time is O((K + 1) * E), which is O(K * E) for large K. It does not depend on n directly, because each flight does a constant amount of work. The copy takes O(n) per round, so the full cost is O((K + 1) * (E + n)). The extra n term is small in practice, but it is there, and a careful answer names it.

Summary

• The limit counts stops, and K stops means at most K + 1 flights, so run K + 1 rounds or levels.
• Each round must read the distances from the start of the round and write into a copy, or one round can chain several flights.
• The rounds method and the level method give the same answers, and the level method can stop early when no price changes.
• Use math.inf for "not reached", so that unreached cities never look cheap.
• A cheap route can use too many stops, so the answer for a limit can differ from the cheapest route overall.