Memoization

Introduction

Some recursive functions solve the same small problem again and again. Memoization stores each answer the first time it is computed and looks it up afterwards.

Some recursive functions are beautifully short and hopelessly slow, because they solve the same small problems again and again. A small change, remembering answers, can turn a program that would take years into one that runs instantly.

How it works

The slow way

calls = 0

def fib(n):
    global calls
    calls += 1
    if n < 2:
        return n
    return fib(n - 1) + fib(n - 2)

print(fib(8), "after", calls, "calls")

fib(8) needs fib(7) and fib(6), but fib(7) also needs fib(6), and so on. The same calls repeat, so computing fib(8) takes 67 calls. Watch the call stack rise and fall over and over. The work roughly doubles with each step of n: fib(40) would take hundreds of millions of calls.

The built-in way

from functools import lru_cache

@lru_cache(maxsize=None)
def fib(n):
    if n < 2:
        return n
    return fib(n - 1) + fib(n - 2)

print(fib(30))

@lru_cache from the functools module adds memoization to any function in one line. fib(30), which would take over a million calls without it, now takes about 30. (The @ line is a decorator: it wraps the function with extra behaviour.)

This idea, solving each smaller problem once and reusing the answer, is the heart of dynamic programming, a major topic in algorithms courses.

Syntax

memo = {}

def f(n):
    if n in memo:
        return memo[n]
    ...
    memo[n] = result
    return result

Example

calls = 0
memo = {}

def fib(n):
    global calls
    calls += 1
    if n in memo:
        return memo[n]
    if n < 2:
        result = n
    else:
        result = fib(n - 1) + fib(n - 2)
    memo[n] = result
    return result

print(fib(8), "after", calls, "calls")

Before computing, the function checks memo. After computing, it stores the result. Watch the memo dictionary fill up in the objects panel. Each value of n is now computed only once, so fib(8) takes 15 calls instead of 67.

Line by line

1. calls counts how many times the function runs, and memo stores the answers.
2. Each call adds 1 to calls, then returns at once if the answer is known.
3. The base case handles n < 2.
4. Otherwise, the two recursive calls build the answer.
5. memo[n] = result stores it before returning.
6. The output shows the answer, 21, after only 15 calls.

Expected result

21 after 15 calls

In the real world

Web servers cache expensive pages so that they do not have to rebuild them for every visitor, and browsers cache files for the same reason. Spell checkers and route planners rely on dynamic programming, which is memoization's close relative.

Common mistakes

• Memoizing a function whose result depends on something besides its arguments, such as the time or a global variable.
• Using a list or dictionary as an argument to an lru_cache function: arguments must be hashable, such as numbers, strings or tuples.
• Forgetting to store the result in the memo, so nothing is ever remembered.

Best practices

• Check the memo before doing any work.
• Make sure that the arguments can be used as dictionary keys.
• Measure how many calls you save, so you know it was worthwhile.

Summary

Memoization stores the answer for each input the first time it is computed. Overlapping subproblems then cost almost nothing. @lru_cache adds this to a function in one line, and the idea underlies dynamic programming.

Practice

Count the ways to climb n stairs taking 1 or 2 steps at a time, with memoization. Check that 5 stairs give 8 ways.

Mark lesson complete