Python problem 25 of 30

Write a memoize decorator

  • Python
  • Hard

The recursive fib below is far too slow for n = 80. Write a decorator memoize that caches results, so the same code returns instantly.

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

fib(80)   # 23416728348467685
Show hint

A decorator is a function that takes a function and returns a new one. Give the new function a dictionary that maps arguments to results.

Show solution
def memoize(fn):
    cache = {}

    def wrapper(*args):
        if args not in cache:
            cache[args] = fn(*args)
        return cache[args]

    return wrapper

@memoize means fib = memoize(fib), so the name fib now refers to wrapper. That matters: the recursive calls inside fib also go through the cache. Each value is computed once, turning exponential time into O(n).

The arguments are tuple keys, so they must be hashable. In real code use functools.cache, and add functools.wraps(fn) to keep the original function's name and docstring.

Want a problem like this every day? See today's problem and keep a streak going.