Remember the Answer

In Recursion you met a function that was correct and useless. ways counted the ways to climb a staircase, and for 30 stairs it made more than two and a half million calls. It was not doing a great deal of work. It was doing a small amount of work a great many times. This module is about the cure, which is almost embarrassingly simple: the first time an answer is worked out, write it down, and the next time it is wanted, look it up.

Before this, you can: write a function that calls itself, and count its calls; look up and store in a dictionary, with a pair as a key; count the work a function does by what each input can cause.

By the end you will be able to find the questions a function asks more than once, keep its answers in a dictionary and look before working, decide what the key has to hold, tell when an answer is safe to remember, and work out what a remembered function costs.

Before you read on, decide what this prints:

asked = {}

def ways(n):
    asked[n] = asked.get(n, 0) + 1
    if n <= 1:
        return 1
    return ways(n - 1) + ways(n - 2)

print(ways(6))
print(asked[6], asked[4], asked[2], asked[1])

13
1 2 5 8

There are 13 ways to climb six stairs. To find that out, the function asked for the answer to four stairs twice, to two stairs five times, and to one stair eight times. Each of those answers was the same every time. It was worked out from nothing every time.

The same question, again and again

ways(6) needs ways(5) and ways(4). But ways(5) needs ways(4) as well, so four stairs is worked out twice. Each of those needs ways(3), and so does ways(5) directly: three times. The smaller the input, the more often it is asked for, and the counts grow in the same pattern as the answers themselves: 1, 1, 2, 3, 5, 8.

Seven different inputs are involved, 0 to 6. Yet 25 calls were made. For 30 stairs there are 31 different inputs and 2,692,537 calls. The whole cost of this function is in repeats.

This is the thing to look for. When a function calls itself more than once, ask: can two different chains of calls arrive at the same input? If they can, the function is repeating itself, and it can be cured. If they cannot, as with a folder of folders where every folder is reached one way only, there is nothing to remember.

Remember each answer

Keep a dictionary of answers. At the top of the function, look in it. If the answer is there, hand it back. If not, work it out as before, and store it before returning.

memo = {}
calls = {"count": 0}

def ways(n):
    calls["count"] = calls["count"] + 1
    if n in memo:
        return memo[n]
    if n <= 1:
        return 1
    answer = ways(n - 1) + ways(n - 2)
    memo[n] = answer
    return answer

print(ways(6), calls["count"])
calls["count"] = 0
memo.clear()
print(ways(30), calls["count"])

13 11
1346269 59

Fifty-nine calls for thirty stairs, where there were more than two and a half million. The function is the same function with three lines added: a look, and a store just before the return. The dictionary is called a memo, as in a note to yourself, and the technique is called memoization.

Notice what did not change. The case that stops is the same. The step is the same. You do not have to understand the problem any better to speed it up: you write the plain function that calls itself, make sure it is right on small inputs, and then add the memo.

Why 59? Each input from 2 to 30 is worked out exactly once, and working it out makes two calls. Most of those calls find their answer already in the memo and return at once.

What is the key?

For ways the key was easy: the function has one input, and the answer depends on it alone. When a function has several inputs, the key must hold everything the answer depends on.

Take routes on a grid of streets, walking only across or down. The number of routes depends on both how far across and how far down there is to go, so the key is the pair.

memo = {}

def routes(across, down):
    if across == 0 or down == 0:
        return 1
    key = (across, down)
    if key in memo:
        return memo[key]
    answer = routes(across - 1, down) + routes(across, down - 1)
    memo[key] = answer
    return answer

print(routes(2, 2), routes(8, 8))
print(len(memo))

6 12870
64

Sixty-four answers are in the memo afterwards: one for each pair with both numbers from 1 to 8. Recursion's version made 25,739 calls for this grid.

Suppose the key had been across alone. Then routes(1, 5) and routes(1, 2) would share one entry, the second would be handed the first one's answer, and the function would be wrong with no error at all. A key that holds too little gives wrong answers. A key that holds more than it needs is only wasteful: answers that could have been shared are worked out separately. When in doubt, put it in the key.

Where you are, not how you got there

A memo hands back an old answer for a new call. That is only right if the two calls really are the same question. So the answer must depend on the inputs in the key, and on nothing else: not on a list that changes between calls, not on a total kept outside, and not on the path taken to get here.

That last one is the usual trap. Suppose you count routes across a grid where some squares may each be used only once. The number of ways onward from a square now depends on which squares you have already used, which is to say on how you got there. Two calls with the same position are not the same question, and a memo keyed on position alone would be wrong.

There are two honest ways out. Put the extra thing in the key, if it is small: a position and a budget left, say. Or restate the problem so that the question at each point looks only forward. Most of the skill in this technique is choosing what "where you are" means so that it says everything that matters and nothing more.

The same reasoning says when a memo must be thrown away. routes above remembers answers for an open grid. Ask it about a grid with a closed crossing, using the same dictionary, and it will hand back answers for the open one. A memo belongs to one problem. If anything the answers depend on changes, start with an empty one.

Counting the work

Without a memo, the number of calls can double at every level, and the only way to find the cost is to count. With a memo there is a rule.

The work is the number of different questions, times the work to answer one of them given the smaller answers.

For the stairs: about n different questions, and each is two look-ups and an addition. So the number of steps is in step with n. For the grid: across × down questions, each a look-up or two and an addition.

That counts steps, taking a dictionary look-up and an addition as one step each. It is the right count for comparing two ways of solving a problem, and it leaves two things out. A look-up is one step on average, as in Lookup in One Pass. And the numbers themselves grow: the ways up sixty stairs is a number of thirteen digits, and adding long numbers takes longer than adding short ones. For answers of the sizes in this module that does not change the picture. A 20 by 20 grid is 400 questions, where the plain function would make over two hundred thousand million calls.

So the first thing to do with any problem of this kind is count the different questions. If there are a thousand, or a million, the repeats are gone and the work is that many steps. If the number of different questions is itself enormous, because the key has to hold too much, no memo will save it.

A memo removes repeated work and nothing else. Every answer is still stored, and the calls still wait inside one another: a chain of questions ten thousand long, each needing the one before, will stop with a RecursionError however well it remembers. A chain that long has to be worked through by a loop.

Where it goes wrong

• The answer is stored under the wrong key, or looked up under one key and stored under another. The function is right the first time and wrong later, or simply never finds anything it stored.
• The key leaves out something the answer depends on. No error: wrong answers, and only for some inputs.
• The memo outlives its problem. Answers from one grid, one price list or one set of rules are handed back for another.
• The store is skipped on one of the paths. A function with two return lines after the look-up needs a store before each, or it remembers only some of its answers and stays slow.
• The function does more than answer. A function that prints, or adds to a list, does that only the first time it meets an input. A remembered function should return a value and do nothing else.

In the lab

You will count how often each question is asked, add a memo and count again, and write the stairs function for a staircase with broken steps. You will decide the key for a grid, count routes round closed crossings, and see a memo give a wrong answer because its key held too little. You will find the best way down a triangle of numbers and decide whether a text can be cut into words. Four assignments follow: the fewest stamps for a postage, a memo with two mistakes, a performer who never plays two nights running, and a message whose spaces were lost.