Memoize it, the Python way!

by jasonyeo

I am currently reading Python Algorithms by Magnus Hetland Lie. It’s an awesome book about coding algorithms in python. You should check it out. One of the things that I thought it was really interesting is the use of decorators when memoizing in python.

In case you don’t know, memoization is the use of caches to retrieve previously computed values to reduce the time complexity of algorithms.

For example, this is the code for to compute the nth fibonacci number:

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

The running time for this piece of code is exponential. The memoized version of the function goes like this:

f = [-1]*n
def pre_fib(n):
    f[0] = 1
    f[1] = 1
def fib(n):
    if f[n] != -1: return f[n]
    return f[n-1] + f[n-2]

This version of the fibonacci function is polynomial. The key lies in the use of the array to retrieve previously computed values instead of computing it again.

Decorate Your Python

In Python, we can transform functions using decorators. It is a design pattern that allows the programmer to add behavior to existing objects. To decorate our fibonacci function, we need to first define what behavior we want to add to the function. This is how we define the decorator:

from functools import wraps

def memo(func):
    cache = {}
    def wrap(*args):
        if args not in cache:
            cache[args] = func(*args)
        return cache[args]
    return wrap

This piece of code simply wraps the function that we want to memoize and calls it if the value we are looking for is not in the cache. If it is in the cache, it simply returns the value from the cache. Here comes the beauty of using decorators to memoize the fibonacci function. IMO, it’s really elegant. See it for yourself:

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

That’s it! It’s really easy. Now simply call the fib function to see it in action!