Memoization : Using Decorators to Speed Up Recursion

One big issue with a lot of recursive algorithms is that they can lead to very inefficient algorithms that often demonstrate exponential time requirements. In these cases a technique called memoization can be really useful. Note that if you are looking for notes on memoization it is spelt the US English way, not as memoisation.

Lets consider a classic combinatorial problem: Suppose you are climbing a set of n stairs taking either 1,2 or 3 steps each time. How many ways are there to climb the stairs?

For low N we may be able to calculate this by hand, however for higher numbers of steps a computer program seems like a wise idea. Lets write a naive recursive program in python to solve this problem

# hopping up steps problem. N-steps, 1, 2 or 3 steps at a time, 
#how many ways can we climb the steps?

# number of steps
steps = 10

def one_step(steps_to_go):
    '''given a number of steps to climb recursively climb either 
    1, 2, or 3 steps once we reach the end of the steps pass 
    totalled step values back up the recursive chain'''
    # try to take one step
    if steps_to_go == 1:
        # upon finding there is only one step to go return 1 
        # as there is only 1 way to climb this
        return 1
    # try to take 2 steps
    elif steps_to_go == 2:
        # upon finding there are 2 steps to go return 2 
        # (2 ways to climb: 2; 1,1)
        return 2
    # try to take 3 steps
    elif steps_to_go == 3:
        # upon finding there are 3 steps to go return 4 
        # (4 ways to climb: 3; 2,1; 1,2; 1,1,1)
        return 4
    else:
        # add all the results of stepping again 
        # (taking each of the 3 different steps)
        return (one_step(steps_to_go-1) + one_step(steps_to_go-2)
                + one_step(steps_to_go-3))

# print our answer
print(one_step(steps))

This works for low step numbers, but rapidly becomes intractable due to the fact that this naive solution is O(3^n) complexity since each call to one_step tends to call 3 more copies of the function. Trying to solve for 100 steps this way would be highly impractical.

However a key feature to recognise is that each branch of this problem will actually be solving the same subproblems (how many ways are there to climb x steps for each x less than n) over and over again. this means that we can attack this problem using memoization. Note that this is not true of divide and conquer algorithms like quicksorts which are composed of subproblems, but not overlapping ones.

What memoization does is to keep a tally of which subproblems are already solved and what the answers to those subproblems are.  A really nice feature of memoization using decorators is that it does not need us to refactor our existing recursive code. We can add memoization as a wrapper around our existing code.

The code for the memoization decorator is very simple. It is the standard wrapper function inside a decorator function and a small amount of logic to check if the wrapped function has a known value for the current argument.

def memoize(func):
    '''Memoizing decorator that checks whether the wrapped function 
    has been previously run with the current argument value. 
    If it has then that value is retrieved and offered instead
    of rerunning the function'''
    memos = {}
    def wrapper(x):
        if x not in memos:
            memos[x] = func(x)
        return memos[x]
    return wrapper

Remember that because of how decorators work, the arguments of wrapper are the arguments of the function being wrapped. So what this code is doing is looking to see if the wrapped recursive function has previously been called with those arguments. If so it retrieves and returns the answer from the  memos dictionary, otherwise, it runs the function call and stores the result in memos before returning it.

This does not just reduce operations to a lookup table though. It actually prunes entire branches off the tree and never runs them where the answer is already know, This means that it reduces run time to O(n) time and also reduces storage requirements to O(n) at the same time. This means that we can attack vastly larger values of n than were possible with naive recursion.

For example we can run the memoized code below to almost instantly get the answer for climbing a 1000 step staircase. This would have been completely impractical with the naive code

# for large numbers we will have to memoize!
steps = 1000

# generic memoizer
def memoize(func):
    '''Memoizing decorator that checks whether the wrapped function 
    has been previously run with the current argument value. 
    If it has then that value is retrieved and offered instead 
    of rerunning the function'''
    memos = {}
    def wrapper(x):
        if x not in memos:
            memos[x] = func(x)
        return memos[x]
    return wrapper

@memoize
def one_step(steps_to_go):
    '''given a number of steps to climb recursively climb either 
    1, 2, or 3 steps once we reach the end of the steps pass 
    totalled step values back up the recursive chain'''
    # try to take one step
    if steps_to_go == 1:
        # upon finding there is only one step to go return 1 
        # as there is only 1 way to climb this
        return 1
    # try to take 2 steps
    elif steps_to_go == 2:
        # upon finding there are 2 steps to go return 2 
        # (2 ways to climb: 2; 1,1)
        return 2 
    # try to take 3 steps
    elif steps_to_go == 3:
        # upon finding there are 3 steps to go return 4 
        # (4 ways to climb: 3; 2,1; 1,2; 1,1,1)
        return 4 
    else:
        # add all the results of stepping again 
        # (taking each of the 3 different steps)
        return (one_step(steps_to_go-1) + one_step(steps_to_go-2)
                + one_step(steps_to_go-3))

# print our answer
one_step(steps)

Memoization is one of the two main strategies in Dynamic Programming. Memoization is a top-down method for storing intermediate results, while tabulation is the bottom up method. Tabulation may have advantages for certain classes of problems, but memoization is exceptionally easy to implement and allows the use of existing code.