How to apply memoization in Python using functools

Memoization is a technique used to optimize an exponential recursive function into a fast polynomial one.

Let’s take the classic problem of Fibonacci numbers. The recursive \(O(2^n)\) function would look like this:

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

In order to obtain an \(O(n)\) time complexity, we want to transform it using memoization. Because \(fib(n)\) is called multiple times, we want to store the result in a cache the first time we compute it, and retrieve it from there on subsequent calls.

def fib(n, cache = {}):
  if n < 2:
    return n
  if n in cache:
    return cache[n]
  cache[n] = fib(n - 1, cache) + fib(n - 2, cache)
  return cache[n]

Functools offers a way to cache the results of a function, considering the parameters are hashable, in a similar way the functional programming languages do.

from functools import lru_cache

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

We can notice that the fib function remained the same, the only thing we added is a decorator @lru_cache with None as parameter, meaning the cache size is unlimited