{"id":"dsa-dp-memoization","slug":"memoization","title":"Memoization","text":"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.\n\nThe 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:\n\n```python\ncalls = 0\n\n\ndef fib(n):\n    global calls\n    calls += 1\n    if n < 2:\n        return n\n    return fib(n - 1) + fib(n - 2)\n\n\nfor n in [10, 20, 25]:\n    calls = 0\n    print(f\"fib({n}) = {fib(n)}, {calls} calls\")\n```\n\n```output\nfib(10) = 55, 177 calls\nfib(20) = 6765, 21891 calls\nfib(25) = 75025, 242785 calls\n```\n\n### Why it's slow\n\n`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.\n\n### A memo\n\n**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:\n\n```python\ncalls = 0\n\n\ndef fib(n, memo):\n    global calls\n    calls += 1\n    if n < 2:\n        return n\n    if n in memo:\n        return memo[n]           # solved before: look it up\n    result = fib(n - 1, memo) + fib(n - 2, memo)\n    memo[n] = result             # store it for next time\n    return result\n\n\nfor n in [10, 20, 25]:\n    calls = 0\n    print(f\"fib({n}) = {fib(n, {})}, {calls} calls\")\nprint(fib(90, {}))\n```\n\n```output\nfib(10) = 55, 19 calls\nfib(20) = 6765, 39 calls\nfib(25) = 75025, 49 calls\n2880067194370816120\n```\n\nEach `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).\n\nSo 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.\n\n### functools.cache\n\nThe 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:\n\n```python\nfrom functools import cache\n\n\n@cache\ndef fib(n):\n    if n < 2:\n        return n\n    return fib(n - 1) + fib(n - 2)\n\n\nprint(fib(90))\nprint(fib.cache_info())\n```\n\n```output\n2880067194370816120\nCacheInfo(hits=88, misses=91, maxsize=None, currsize=91)\n```\n\nThe 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.\n\n### The recursion limit\n\nA 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:\n\n```python\ndef fib(n, memo):\n    if n < 2:\n        return n\n    if n not in memo:\n        memo[n] = fib(n - 1, memo) + fib(n - 2, memo)\n    return memo[n]\n\n\nprint(fib(500, {}) % 1000)\nprint(fib(5000, {}) % 1000)\n```\n\n```output\n125\n```\n\n```crash\nTraceback (most recent call last):\n  File \"main.py\", line 10, in <module>\n    print(fib(5000, {}) % 1000)\n          ~~~^^^^^^^^^^\n  File \"main.py\", line 5, in fib\n    memo[n] = fib(n - 1, memo) + fib(n - 2, memo)\n              ~~~^^^^^^^^^^^^^\n  File \"main.py\", line 5, in fib\n    memo[n] = fib(n - 1, memo) + fib(n - 2, memo)\n              ~~~^^^^^^^^^^^^^\n  File \"main.py\", line 5, in fib\n    memo[n] = fib(n - 1, memo) + fib(n - 2, memo)\n              ~~~^^^^^^^^^^^^^\n  [Previous line repeated 996 more times]\nRecursionError: maximum recursion depth exceeded\n```\n\n`@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.\n\n### What it costs\n\nA 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.\n","task":"`choose(n, k, memo)` counts the ways to choose k of n items. It's the take-or-leave choice of the backtracking module: either the last item is taken, and the other k - 1 come from the first n - 1 items, or it's left, and all k come from the first n - 1. Without a memo, `choose(40, 20, {})` makes over 275 billion calls. Fill in the blanks: the memo's key, which must hold both numbers that change, and the check that the answer is already in the memo. The program should print:\n\n```\n10\n137846528820\n```\n","kind":"fill","seed":"def choose(n, k, memo):\n    if k == 0 or k == n:\n        return 1\n    key = [[(n, k)‖n, k]]\n    if [[key in memo]]:\n        return memo[key]\n    memo[key] = choose(n - 1, k - 1, memo) + choose(n - 1, k, memo)\n    return memo[key]\n\n\nprint(choose(5, 2, {}))\nprint(choose(40, 20, {}))\n","solution":"def choose(n, k, memo):\n    if k == 0 or k == n:\n        return 1\n    key = (n, k)\n    if key in memo:\n        return memo[key]\n    memo[key] = choose(n - 1, k - 1, memo) + choose(n - 1, k, memo)\n    return memo[key]\n\n\nprint(choose(5, 2, {}))\nprint(choose(40, 20, {}))\n","hints":["The answer depends on both `n` and `k`, so the key must hold both of them: a tuple.","The key is `(n, k)`, and the check asks whether that key is already in the dict: `key in memo`."],"tests":[{"name":"Output","stdin":"","expect":"10\n137846528820","hidden":false}],"require":[],"forbid":[],"more":[{"task":"What does this program print? Type each line it prints.\n","kind":"predict","seed":"from functools import cache\n\n\n@cache\ndef f(n):\n    print(\"solving\", n)\n    if n < 2:\n        return n\n    return f(n - 1) + f(n - 2)\n\n\nprint(f(4))\nprint(f(5))\n","solution":"from functools import cache\n\n\n@cache\ndef f(n):\n    print(\"solving\", n)\n    if n < 2:\n        return n\n    return f(n - 1) + f(n - 2)\n\n\nprint(f(4))\nprint(f(5))\n","hints":["`f(4)` calls `f(3)` first, which calls `f(2)` first, and so on down to `f(1)`. Each of them prints when it starts.","When `f(2)` then asks for `f(0)`, that one is new too. But when `f(3)` asks for `f(1)`, and `f(4)` for `f(2)`, the answers come from the cache: nothing is printed.","`f(5)` is new, so it prints. It needs `f(4)` and `f(3)`, which are both in the cache already."],"tests":[{"name":"Output","stdin":"","expect":"solving 4\nsolving 3\nsolving 2\nsolving 1\nsolving 0\n3\nsolving 5\n5","hidden":false}],"require":[],"forbid":[],"lines":["solving 4","solving 3","solving 2","solving 1","solving 0","3","solving 5","5"]},{"task":"Write `rolls(dice, faces, total)`: in how many ways can `dice` dice, each with the faces 1 to `faces`, show numbers that add up to `total`? The dice are different dice, so a 1 on the first and a 2 on the second is another way than a 2 on the first and a 1 on the second.\n\nThe last die shows some face f from 1 to `faces`, and the other dice must then make `total - f`: add up those counts. With no dice left, there's exactly one way to make a total of 0, and none to make anything else. Plain recursion tries every one of the `faces ** dice` rolls, and the tests go up to 100 dice, so memoize.\n\nFor example, `rolls(2, 6, 7)` returns 6 (1 + 6, 2 + 5, 3 + 4, 4 + 3, 5 + 2 and 6 + 1):\n\n```\n6\n```\n","kind":"code","seed":"def rolls(dice, faces, total):\n    pass\n\n\nif __name__ == \"__main__\":\n    print(rolls(2, 6, 7))\n","solution":"from functools import cache\n\n\n@cache\ndef rolls(dice, faces, total):\n    if dice == 0:\n        return 1 if total == 0 else 0\n    ways = 0\n    for face in range(1, faces + 1):\n        if face <= total:\n            ways += rolls(dice - 1, faces, total - face)\n    return ways\n\n\nif __name__ == \"__main__\":\n    print(rolls(2, 6, 7))\n","hints":["Start with the plain recursion. The base case is `dice == 0`: return 1 when `total == 0`, and 0 otherwise. Otherwise loop over the faces and add up `rolls(dice - 1, faces, total - face)`.","Then memoize it: `from functools import cache` and `@cache` above the function. Or keep a dict outside the function (or let `rolls` call a helper that takes the memo along), with a key that holds every argument that changes. The tests call `rolls` with its three arguments only, so don't give it a fourth parameter for the memo.","A dict that is shared by every call of `rolls` needs the faces in its key as well: `rolls(2, 6, 7)` and `rolls(2, 4, 7)` ask different questions."],"tests":[{"name":"rolls(2, 6, 7)","stdin":"","call":"rolls(2, 6, 7)\n","expect":"6","hidden":false,"returns":true},{"name":"rolls(1, 6, 7)","stdin":"","call":"rolls(1, 6, 7)\n","expect":"0","hidden":true,"returns":true},{"name":"rolls(3, 6, 10)","stdin":"","call":"rolls(3, 6, 10)\n","expect":"27","hidden":true,"returns":true},{"name":"rolls(0, 6, 0)","stdin":"","call":"rolls(0, 6, 0)\n","expect":"1","hidden":true,"returns":true},{"name":"rolls(3, 6, 2)","stdin":"","call":"rolls(3, 6, 2)\n","expect":"0","hidden":true,"returns":true},{"name":"print(rolls(2, 6, 7), rolls(2, 4, 7), rolls(4, 2, 6))","stdin":"","call":"print(rolls(2, 6, 7), rolls(2, 4, 7), rolls(4, 2, 6))\n","expect":"6 2 6","hidden":true},{"name":"rolls(30, 6, 105)","stdin":"","call":"rolls(30, 6, 105)\n","expect":"9378595792117360310832","hidden":true,"big":true,"returns":true},{"name":"rolls(100, 6, 350) % 1_000_000_007","stdin":"","call":"rolls(100, 6, 350) % 1_000_000_007\n","expect":"850903066","hidden":true,"big":true,"returns":true},{"name":"rolls(40, 10, 220) % 1_000_000_007","stdin":"","call":"rolls(40, 10, 220) % 1_000_000_007\n","expect":"388529127","hidden":true,"big":true,"returns":true}],"require":[],"forbid":[]}]}