Memoised recursion in Python

Why this matters

The last lesson drew the call tree for naive Fibonacci and found that the same calls appear many times. fib(2) appeared twice in the tree for fib(4), and fib(30) makes about 2.7 million calls, most of them repeats. A function that does the same work again and again is slow for no good reason. Its answer for a given input never changes, so it could have been computed once and reused.

This lesson adds that reuse. You will keep a record of every answer the function has already found, and check that record before doing any new work. The technique is called memoisation. The word comes from "memo," a note you write to remember something. It is the first step into dynamic programming, the method of solving a large problem by solving its smaller parts once and reusing the answers. Dynamic programming has its own module later in this course. This lesson covers only the top-down form, where the recursion stays and the answers are remembered.

The idea from the first principle

Imagine a student who must answer the same question on a test several times. The first time, the student works it out and writes the answer in the margin. The second time, the student reads the margin and does not work it out again. The margin note is a memo.

A recursive function can do the same. It needs a place to store answers, and a way to look them up. In Python the usual choice is a dict. A dict stores pairs of a key and a value, and it finds a value from its key quickly. A dict[int, int] stores an int answer for each int input. The key and value types are written in the type hint, and Python does not check them when the program runs, so the hint is a note for readers and tools.

Every memoised function follows the same three steps:

1. Check the dict. If the answer for this input is already there, return it at once.
2. Compute the answer with the recursive calls, as before.
3. Store the answer in the dict, then return it.

The base cases still come first, because they need no storage. The check in step 1 is what turns an exponential tree into a short list of new work.

Worked example: Fibonacci with a dict

Here is the naive Fibonacci function from the last lesson, with the three steps added. The dict is a global variable, so every call shares the same one. The counter calls records how many times the function runs, so you can see the effect of the dict:

Press Run. You should see:

fib(10) = 55, calls = 19
fib(50) = 12586269025, calls = 99, memo size = 49
fib(92) = 7540113804746346429, calls = 183
fib(93) = 12200160415121876738, calls = 185

Compare the first line with the last lesson. There, fib(10) made 177 calls. Now it makes 19. The answer is the same, 55, but the work is much smaller.

Here is why the count is 19. The function computes a real answer for each n from 2 up to 10, which is 9 values. Each of those computed values makes two calls, one for n - 1 and one for n - 2. The first call, fib(10), is counted once on its own. So the total is 1 + 2 * 9, which is 19. The calls for n - 1 and n - 2 that hit the dict or the base case still count as calls, but they do not start any new work.

The same pattern gives 99 calls for fib(50). The computed values run from 2 to 50, which is 49 values, and the dict holds 49 entries. The total is 1 + 2 * 49, which is 99. The count now grows in a straight line with n, not exponentially. For fib(50), the naive method would make about 40.7 billion calls, since 2 * fib(51) - 1 is 2 * 20365011074 - 1, which is 40730022147. The memoised method makes 99.

Now follow the first steps by hand for fib(4), with an empty dict. The call fib(4) looks in the dict and finds nothing. It calls fib(3), which also finds nothing and calls fib(2). The call fib(2) finds nothing, calls fib(1) and fib(0), which return 1 and 0 directly, and stores 1 for key 2. Back in fib(3), the second call is fib(1), which is a base case, so it returns 1. The value 2 is stored for key 3. Back in fib(4), the second call is fib(2). This time the dict has key 2, so the call returns 1 at once, without calling anything. The value 3 is stored for key 4. The second fib(2) in the naive tree is the saved work.

The last two lines show that Python integers are exact. A Python int grows to hold any size, so fib(93) is printed correctly. The dict saves work, but it does not limit the size of the numbers. The cost of a large answer is memory and time for the arithmetic, not a wrong result.

Second example: paths on a grid

The first example had one input number. The second example has two. Imagine a grid of streets. You start at the top left corner and want to reach a point that is row steps down and col steps right. You can move only right or down. How many different routes are there?

A route to a point has a simple rule. The last move was either a step down from the point above it, or a step right from the point to its left. So the number of routes to a point is the number of routes to the point above plus the number of routes to the point to its left. On the top row and the left column there is only one route, the straight line.

The dict now needs two numbers as its key. One way is to build a string such as "2,3", but that is easy to get wrong. A better way is a tuple, which is a fixed group of values written in parentheses, such as (2, 3). Tuples can be dict keys, and two tuples with the same values are equal, so the dict treats them as the same key.

Press Run. You should see:

paths(2, 2) = 6, calls = 9, memo size = 4
paths(3, 3) = 20, calls = 19, memo size = 9
paths(5, 5) = 252, calls = 51, memo size = 25
paths(8, 8) = 12870, calls = 129, memo size = 64
paths(10, 10) = 184756, calls = 201, memo size = 100
paths(2, 3) = 10
paths(3, 2) = 10
paths(16, 16) = 601080390, calls = 513

Look at the memo size column. For a 10 by 10 grid the dict holds 100 entries, one for each point that is not on the top row or the left column. Each of those points is worked out once. The call count follows the same rule as before: 1 + 2 * 100, which is 201.

The table also shows the value of the method. A plain recursive version of this function, with the dict lines removed, makes 369511 calls for the 10 by 10 grid. You can check that by deleting the dict lines. The memoised version makes 201. For the 16 by 16 grid the memoised version makes 513 calls and returns 601080390, a number that would take far too many calls to reach without the dict.

Check paths(2, 3) by hand. Every route is a sequence of two steps down and three steps right, in some order. The number of orders is the number of ways to choose which two of the five moves are downs, which is 10. The program agrees.

The key here is the tuple. If the key were a mutable list, such as [row, col], Python would refuse to use it. Lists can change, so they cannot be dict keys. The fourth common mistake below shows that error.

Third example: fewest coins

The third example is a problem with a choice at each step. You need to make an amount of money with coins, and you want the fewest coins. The coins here are 1, 5, 10, and 25 cents. For an amount n, you pick one coin c and then make the remaining n - c with as few coins as possible. You try every coin and keep the best result. If the amount is 0, no coins are needed. If no combination of coins makes the amount, the answer is a large sentinel value.

A sentinel is a special value used to mean "not a real answer." Here it is math.inf, which means infinity. It is larger than any number, so min never picks it over a real answer. The code never adds one to it, so the sentinel never grows into a wrong count. Python's math.inf is a float, so the function returns a float for an impossible amount. The program never reaches that case, because the 1-cent coin makes every amount.

Press Run. You should see:

no memo: amount 30 needs 2 coins, calls = 16068
no memo: amount 40 needs 3 coins, calls = 343065
memo: amount 30 needs 2 coins, calls = 84, memo size = 30
memo: amount 40 needs 3 coins, calls = 124, memo size = 40
memo: amount 63 needs 6 coins, calls = 216, memo size = 63
memo: amount 99 needs 9 coins, calls = 360, memo size = 99

The amount 30 is 25 + 5, so 2 coins. The amount 40 is 25 + 10 + 5, so 3 coins. The amount 63 is 25 + 25 + 10 + 1 + 1 + 1, so 6 coins. The amount 99 is 25 + 25 + 25 + 10 + 10 + 1 + 1 + 1 + 1, so 9 coins. The program finds all four.

The difference in calls is large. For amount 40 the plain function makes 343065 calls, and the memoised one makes 124. The plain version asks for the same amounts again and again, the same way the naive Fibonacci did. The dict stores one answer per amount. It has one entry for each amount from 1 to the target, and that is why the memo size equals the amount.

Notice what the key is. The dict uses only amount as its key, because the coin list never changes. If the coins were a parameter of the function, the key would need the coin list as well, or at least the amount and a position in the list. Before you choose a key, list every input that changes the answer. The key must include all of them and nothing more.

Python has a shorter way to write the memo for this kind of function. The @cache decorator from the second lesson stores answers for any function whose arguments are hashable. The explicit dict in this lesson is worth writing at least once, because it shows exactly what is stored and when. Use @cache in your own code when it fits.

Common mistakes

Students who first write memoised functions hit the same few problems. Each one below shows what goes wrong, so you can spot it in your own code.

Mistake 1: checking the dict after the recursion

The three steps have an order. If you compute the answer first and check the dict afterwards, the recursion runs every time, and the dict never saves any work. The check must come before the recursive calls. The output of the first example shows the effect: with the check in the wrong place, fib(10) would still make 177 calls.

Mistake 2: a key that leaves out an input

The key must include every input that changes the answer. The next program is a version of the grid function that uses only the row as its key. It looks like it works, because the function still runs and returns a number:

Press Run. You should see:

buggy paths(2, 3) = 13

The correct answer is 10, as the second example showed. The function returned 13 because the entry for row 1 was saved for one column and then reused for a different column. There is no error message. This is the most dangerous mistake in this lesson, because the output looks like a normal number. To find it, test the function against a small case that you can count by hand, such as paths(2, 3), before you trust a large result. The fix is the tuple key from the second example.

Mistake 3: a dict that holds old answers

The dict is a global variable, so it keeps its entries between calls to the function. If you change the rules of the problem, for example the list of coins, the old entries are still in the dict, and the function returns answers from the old rules. The safe habit is to call memo.clear() before each test case, as the examples in this lesson do. In a real program, the dict should belong to one problem, and a new problem should get a new dict. A default argument such as memo: dict = {} in the function header is a trap for the same reason, because Python creates that dict once and every call shares it.

Mistake 4: using a list as a key

A list cannot be a dict key, because a list can change and a key must stay the same. Python stops the program with an error when you try:

The last line of the output is:

TypeError: cannot use 'list' as a dict key (unhashable type: 'list')

The message says the value is unhashable. A hash is a number that Python computes from a value to find it in a dict quickly. Lists have no hash, because they can change. Tuples and numbers have hashes. The fix is to change [2, 3] to the tuple (2, 3). A similar error, TypeError: unhashable type: 'list', appears when you use @cache on a function that takes a list as an argument.

Mistake 5: a huge answer that costs memory

Python integers never overflow, so a large answer is exact. The cost is memory and time for the arithmetic. A memo that stores very large numbers, one for each input, can use a lot of memory. Check the size of the answers before you store them all, and clear the dict between problems that do not share answers.

Practice

Exercise 1

The grid function makes 201 calls for paths(10, 10). Explain where that number comes from, using the size of the dict.

Show the answer

The dict holds one entry for each point from row 1 to 10 and column 1 to 10, which is 10 times 10, or 100 entries. Each of those 100 computed points makes two calls, one for the point above and one for the point to its left. The first call, paths(10, 10), is counted once on its own. The total is 1 + 2 * 100, which is 201.

Exercise 2

The output of the Fibonacci program shows that fib(93) is 12200160415121876738. Explain why this is a correct answer in Python, and what would happen to the same value in a 64-bit integer.

Show the answer

Python int values grow to hold any size, so the sum is exact. A 64-bit integer holds numbers only up to 9223372036854775807. The value 12200160415121876738 is larger than that, so a 64-bit integer wraps around and stores a wrong negative number, -6246583658587674878. Python has no such limit, so the printed value is right.

Exercise 3

The buggy grid function uses only row as its key. Fix it so that buggy_paths(2, 3) returns 10. Describe the change in one sentence.

Show the answer

Change the key from row to the tuple (row, col), so that the key holds both numbers. Two different points then never share a saved answer. The result for paths(2, 3) is then 10. The code is key = (row, col), then if key in memo:, then memo[key] = answer.

Exercise 4

The function below uses functools.cache and works, but the same idea can be written with a plain dict and no decorator. Rewrite the fib function so that it uses the check, compute, and store pattern with a dict, and still returns the same answers.

from functools import cache
 
 
@cache
def fib(n: int) -> int:
    if n <= 1:
        return n
    return fib(n - 1) + fib(n - 2)

Show the answer

memo: dict[int, int] = {}
 
 
def fib(n: int) -> int:
    if n <= 1:
        return n
    if n in memo:
        return memo[n]
    answer = fib(n - 1) + fib(n - 2)
    memo[n] = answer
    return answer
 
 
print(fib(50))  # 12586269025

The check n in memo comes before the recursive calls, and the store comes after them. The function returns the same answers as the @cache version. The dict is visible, so you can print it or clear it, which the decorator does not let you do directly.

Exercise 5

For the coin program, the memo size for min_coins(63) is 63. Explain why it is 63 and not 64, and find the number of calls the program makes.

Show the answer

The dict stores one entry for each amount from 1 to 63. The amount 0 is handled by the base case before the dict is checked, so it is never stored. That gives 63 entries. The program prints 216 calls for amount 63, as the output above shows. That count includes every call, including the calls that stop at amount 0 and the calls that find an amount already in the dict. Each of those returns at once, but each still adds one to the counter.

Summary

• Memoisation stores the answer to each input in a dict, and checks the dict before doing new work.
• The three steps are: check the dict, compute with the recursive calls, and store the answer before returning it.
• The dict key must include every input that changes the answer, and a tuple makes a good key for two or more numbers.
• A memoised Fibonacci makes 2 * n - 1 calls, so fib(50) makes 99 calls where the naive version makes about 40.7 billion.
• A list cannot be a dict key, because it can change. Use a tuple instead.
• Python integers never overflow, so a memo returns exact answers. The cost of a big memo is memory.