Memoization: remembering answers

The problem: the same work, again and again

The Fibonacci numbers start 0, 1 and each next number is the sum of the two before: 0, 1, 1, 2, 3, 5, 8, 13, … The recursive definition fib(n) = fib(n − 1) + fib(n − 2) is easy to write, but look at the calls for fib(5): it calls fib(4) and fib(3), and fib(4) calls fib(3) again. fib(2) is computed three times. The tree doubles at every level, so the time is O(2ⁿ) (Module 1 showed this recursion tree).

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

print(fib(30), calls)   # Output: 832040 2692537

Overlapping subproblems

A subproblem is a smaller version of the problem, such as fib(3) inside fib(5). The subproblems overlap when the same one is needed many times. Here there are only n + 1 different subproblems (fib(0) to fib(n)), but the naive code solves them millions of times. Compare this with merge sort in Topic 2: its halves never overlap, so remembering would not help there.

Memoization with a dictionary

Memoization means remembering: before computing, look the input up in a cache (a dictionary from input to answer). If it is there, return it. If not, compute it, store it, then return it. Each fib(k) is now computed once, so there are about 2n calls in total, each O(1): O(n) time. The cache holds n + 1 entries and the recursion is n deep, so the space is O(n).

def fib_memo(n, memo=None):
    if memo is None:
        memo = {}
    if n < 2:
        return n
    if n in memo:
        return memo[n]
    memo[n] = fib_memo(n - 1, memo) + fib_memo(n - 2, memo)
    return memo[n]

print(fib_memo(90))   # Output: 2880067194370816120

functools.cache and lru_cache

Python can add the cache for you. Put @functools.cache (Python 3.9 and later) or @functools.lru_cache(maxsize=None) above the function, and every result is stored by its arguments. lru_cache(maxsize=128) keeps only the 128 most recently used results, which saves memory but may recompute. The arguments must be hashable: numbers, strings and tuples work, lists do not (TypeError: unhashable type: 'list'). Call f.cache_clear() to empty the cache, for example between test cases.

from functools import cache

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

print(climb(10))   # Output: 89
print(climb(45))   # Output: 1836311903

Climbing stairs and grid paths

Climbing stairs: you can take 1 or 2 steps at a time; how many ways reach step n? The last move was either 1 step (from n − 1) or 2 steps (from n − 2), so ways(n) = ways(n − 1) + ways(n − 2), with ways(0) = ways(1) = 1. It is Fibonacci in disguise. Grid paths from Topic 1 work the same way: paths(r, c) = paths(r − 1, c) + paths(r, c − 1). Without a memo, a 16 × 16 grid makes hundreds of millions of calls; with one, there are only 256 different (r, c) cells, so it is O(rows × cols). A blocked cell simply returns 0 paths.

Bridge to dynamic programming

Dynamic programming (DP) is the name for solving a problem from the answers of overlapping subproblems. Memoization is the top-down way: write the natural recursion and add a cache. The bottom-up way fills a table in a loop from the smallest subproblem upwards, which avoids the recursion limit and often needs less memory (Fibonacci needs only the last two values). Module 11 is all about DP; the habit to build now is to ask "do the same subproblems repeat?" whenever recursion is slow.

Common mistakes

• Using a mutable default argument, def f(n, memo={}): the same dictionary is shared between ALL calls, even separate test cases with different inputs. Fix: memo=None and create it inside, or use @cache.
• Memoizing on part of the state only, for example caching by r when the answer depends on (r, c). Fix: the cache key must include every argument that changes the answer.
• Passing a list to an @cache function. Fix: pass a tuple, or keep the grid outside and cache on indexes only.
• Expecting memoization to remove the recursion limit: fib_memo(5000) still goes 5000 deep and raises RecursionError. Fix: go bottom-up with a loop for large n.
• Memoizing when subproblems do not repeat (merge sort, permutations). Fix: it only costs memory there; check for overlap first.