Memoization for Recursive Solutions

The Problem with Brute-Force Recursion

Let's quickly recall the recursion tree for finding the number of ways to climb 5 stairs.

Every time our function climbStairs(3) is called, it re-computes the entire branch below it. This redundant work is the source of the inefficiency. The solution seems simple: what if our function could remember the answers it has already found?

Memoization: Remembering Your Results

This is precisely what memoization does. It's an optimization technique where you cache, or store, the results of expensive function calls and return the cached result when the same inputs occur again.

To understand the concept more formally, please read the first few paragraphs of the following article.

The core idea is to trade a bit of memory (to store the results) for a significant gain in speed. When our function needs a result, it first checks its memory (the cache). If the answer is there, it uses it. If not, it does the calculation, stores the answer in the cache, and then returns it.

Implementing Memoization in Python

Let's translate this idea into code. We'll modify the recursive function we developed in the previous lesson. The logic follows a simple, repeatable pattern.

The following video provides an excellent walkthrough of this process using the Fibonacci problem, which has the exact same recursive structure as our Climbing Stairs problem.

As the video demonstrates, the process can be broken down into these steps:

1. Create a cache: A Python dictionary is a perfect choice. It will store our computed results, mapping the input n to the result ways(n).
2. Check the cache: At the very beginning of the recursive function, check if the result for the current input n is already in the dictionary. If it is, return it immediately.
3. Compute and store: If the result is not in the cache, perform the recursive calculation as before. Then, just before returning, store the newly computed result in the cache.

Let's apply this to our climbStairs function.

Original Brute-Force Code (from last lesson):

def climbStairs_recursive(n):



    # Base cases from our previous lesson:
    # ways(0) = 1 (do nothing)
    # ways(1) = 1 (take one step)
    if n == 0:
        return 1
    if n == 1:
        return 1
    



    # Recurrence relation
    return climbStairs_recursive(n - 1) + climbStairs_recursive(n - 2)

Memoized (Top-Down DP) Version:

# We use a helper function to handle the recursion with the memo cache
def climbStairs_helper(n, memo):



    # Step 2: Check if the result is already in the cache
    if n in memo:
        return memo[n]




    # Base cases
    if n == 0:
        return 1
    if n == 1:
        return 1
    



    # Step 3: If not in cache, compute it recursively
    result = climbStairs_helper(n - 1, memo) + climbStairs_helper(n - 2, memo)
    



    # Store the result in the cache before returning
    memo[n] = result
    return result

def climbStairs(n):



    # Step 1: Create the cache (a dictionary)
    memo = {}
    return climbStairs_helper(n, memo)




# Let's test it!
print(f"Ways to climb 5 stairs: {climbStairs(5)}") # Output: 8
print(f"Ways to climb 40 stairs: {climbStairs(40)}") # This now computes almost instantly!

Notice that the core recursive logic climbStairs_helper(n - 1, memo) + climbStairs_helper(n - 2, memo) remains unchanged. We've simply wrapped it with logic to manage our cache.

For another perspective, watch this short segment that directly implements memoization for the Climbing Stairs problem.

Analyzing the Improvement

What have we gained? Let's analyze the time complexity.

With our memoized solution, the recursive function climbStairs_helper is now called for each number from n down to 0 exactly once. The first time it's called for a number k, it does the computation and stores the result. Every subsequent call for k is an O(1) dictionary lookup.

• Time Complexity: We went from exponential O(2^n) to linear O(n). This is a monumental improvement. Calculating climbStairs(40) went from taking minutes (or more) to milliseconds.
• Space Complexity: We are now using extra memory. The recursion can go n levels deep on the call stack, and our memo dictionary will store up to n key-value pairs. Therefore, the space complexity is O(n). This is the trade-off: we use more space to save a lot of time.

A Pythonic Shortcut: @lru_cache

This pattern of creating and managing a cache for a recursive function is so common that Python has a built-in tool to do it for you, called a decorator.

Using this decorator, our entire solution becomes incredibly clean:

from functools import lru_cache

@lru_cache(maxsize=None) # The decorator automatically handles the caching
def climbStairs_pythonic(n):



    # Just the original recursive logic!
    if n == 0:
        return 1
    if n == 1:
        return 1
    return climbStairs_pythonic(n - 1) + climbStairs_pythonic(n - 2)




# Test it
print(f"Ways to climb 40 stairs (pythonic): {climbStairs_pythonic(40)}")

In a technical interview, being able to quickly write this version demonstrates strong Python knowledge. However, it is crucial that you can also explain and implement the manual cache method to show you understand the underlying mechanism.

Conclusion

In this lesson, you applied your first major dynamic programming optimization. By adding a cache to our brute-force recursive solution, we eliminated redundant computations and drastically improved its performance.

Key Takeaways:

• Memoization is a technique to speed up functions by storing the results of previous calls.
• It is the core of the top-down dynamic programming approach.
• The implementation involves three steps: create a cache, check the cache at the start of the function, and store the result in the cache before returning.
• Memoization reduces the time complexity of Fibonacci-style recursive problems from O(2^n) to O(n) at the cost of O(n) space complexity.
• In Python, the @lru_cache decorator provides a convenient way to apply memoization automatically.

Preview of the Next Lesson:

We have now mastered the top-down, recursive approach to dynamic programming. But what if we could solve the problem without recursion at all? In the next lesson, we will explore the bottom-up approach, also known as tabulation. We will see how to solve the same problem by starting with the base cases and iteratively building up our solution, which can sometimes be even more efficient.