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