Memoization

The backtracking module tried every choice, and the number of choices grew exponentially. Often, though, an exponential search asks the same smaller questions again and again. Dynamic programming (DP) answers each different question once, stores the answer, and builds the big answer from the stored small ones.

The classic example is the Fibonacci numbers. Start from 0 and 1, and get every next number by adding up the last two: 0, 1, 1, 2, 3, 5, 8, 13 and so on. You wrote them with recursion in the recursion module. This version counts its calls:

calls = 0


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


for n in [10, 20, 25]:
    calls = 0
    print(f"fib({n}) = {fib(n)}, {calls} calls")

fib(10) = 55, 177 calls
fib(20) = 6765, 21891 calls
fib(25) = 75025, 242785 calls

Why it's slow

fib(25) calls fib(24) and fib(23). But fib(24) calls fib(23) as well, so fib(23) is worked out twice, fib(22) three times, fib(21) five times: the repeats grow like the Fibonacci numbers themselves. Each step of n multiplies the calls by about 1.6, so the time is exponential, about O(1.6ⁿ), and fib(50) would make over 40 billion calls. Yet fib(25) only ever asks 26 different questions, fib(0) up to fib(25): all the rest is repetition.

A memo

Memoization keeps a memo: a dict from each question to its answer. Before working anything out, the function looks in the memo, and after working it out, it stores the answer there:

calls = 0


def fib(n, memo):
    global calls
    calls += 1
    if n < 2:
        return n
    if n in memo:
        return memo[n]           # solved before: look it up
    result = fib(n - 1, memo) + fib(n - 2, memo)
    memo[n] = result             # store it for next time
    return result


for n in [10, 20, 25]:
    calls = 0
    print(f"fib({n}) = {fib(n, {})}, {calls} calls")
print(fib(90, {}))

fib(10) = 55, 19 calls
fib(20) = 6765, 39 calls
fib(25) = 75025, 49 calls
2880067194370816120

Each fib(k) is worked out once. After that, asking again costs one dict lookup, O(1) on average. The answer may be reused because it depends only on the arguments: fib(23) is the same number however you got there. That's what DP needs: the plain recursion asks the same smaller questions many times (overlapping subproblems), and the big answer is built from the answers to smaller questions of the same kind (called optimal substructure when the answer is a best one).

So memoizing a recursive function takes three small changes to it: a memo that every call can reach, a lookup before any work is done, and a store just before the answer is returned. The memo's key must hold every argument that changes. For a function of two numbers, use a tuple, memo[(row, col)]: a tuple can be a dict key, a list can't.

functools.cache

The hash tables module showed functools.lru_cache. Its sibling functools.cache never forgets anything, so it is memoization in one line: put @cache above the function, and Python keeps the dict for you, keyed by all the arguments together:

from functools import cache


@cache
def fib(n):
    if n < 2:
        return n
    return fib(n - 1) + fib(n - 2)


print(fib(90))
print(fib.cache_info())

2880067194370816120
CacheInfo(hits=88, misses=91, maxsize=None, currsize=91)

The 91 misses are the 91 different questions, fib(0) to fib(90), each worked out once, and the 88 hits are repeats answered from the cache. The arguments become dict keys, so they must be hashable: numbers, strings and tuples work, while a list stops with TypeError: unhashable type: 'list'. Pass a tuple or an index instead.

The recursion limit

A memo makes the calls fewer, not shallower. fib(5000) still calls fib(4999), which calls fib(4998), and so on: 5,000 calls waiting at once, and Python stops at about 1,000:

def fib(n, memo):
    if n < 2:
        return n
    if n not in memo:
        memo[n] = fib(n - 1, memo) + fib(n - 2, memo)
    return memo[n]


print(fib(500, {}) % 1000)
print(fib(5000, {}) % 1000)

125

Traceback (most recent call last):
  File "main.py", line 10, in <module>
    print(fib(5000, {}) % 1000)
          ~~~^^^^^^^^^^
  File "main.py", line 5, in fib
    memo[n] = fib(n - 1, memo) + fib(n - 2, memo)
              ~~~^^^^^^^^^^^^^
  File "main.py", line 5, in fib
    memo[n] = fib(n - 1, memo) + fib(n - 2, memo)
              ~~~^^^^^^^^^^^^^
  File "main.py", line 5, in fib
    memo[n] = fib(n - 1, memo) + fib(n - 2, memo)
              ~~~^^^^^^^^^^^^^
  [Previous line repeated 996 more times]
RecursionError: maximum recursion depth exceeded

@cache doesn't get around this: its calls wait on each other the same way. Each one also passes through the cache's own code, written in C, so in the browser, whose stack is smaller, a deep cached recursion can even crash Python outright before it reaches the limit. So keep memoized recursion to small depths (100 dice, a grid of 20 rows), and for big inputs fill the answers in without recursion, from the smallest question up: that's the next step. Don't raise the limit with sys.setrecursionlimit either: every waiting call still takes memory, and the browser's Python can crash outright. Interviewers like to hear you say this.

What it costs

A memoized function costs the number of different questions times the work for each one. fib has n + 1 questions and does O(1) work for each, so it takes O(n) time. The memo holds one answer per question, O(n) memory, and the waiting calls take O(n) more.