如何通过自定义装饰器模式为Python递归爬楼梯算法实现记忆化优化

Great question—building your own memoize decorator is a fantastic way to wrap your head around closures, decorators, and how caching works under the hood. Let's break this down step by step, fixing the scope issue and keeping your core logic clean.

Why Your Initial Attempts Might Have Failed #

When you first tried writing a decorator, you likely ran into two common pitfalls:

• NameError or resetting cache: If you defined the memo dictionary inside the wrapper function, it would get recreated every time the function is called (including recursive calls), so no values stayed cached.
• Shared, unencapsulated cache: Using a global memo dictionary would let all decorated functions share the same cache, leading to bugs and violating basic encapsulation principles.

The fix relies on closures—we’ll create a dedicated cache for each decorated function and store it in the outer scope of the wrapper, so it persists across all calls (including recursive ones).

Step 1: Build the @memoize Decorator #

Here’s a clean, encapsulated implementation that solves the scope problem:

def memoize(func):
    # Create a unique cache for the decorated function (lives in the outer scope)
    memo = {}
    
    def wrapper(n):
        # Check if we've already computed this value
        if n not in memo:
            # If not, run the original function and store the result
            memo[n] = func(n)
        # Return the cached or newly computed value
        return memo[n]
    
    # Return the wrapper, which retains access to `memo` via closure
    return wrapper

Step 2: Apply It to Your count_stairs Function #

Now you can keep your core logic completely untouched—no caching code cluttering the problem-solving part:

@memoize
def count_stairs(n):
    # Base cases (unchanged)
    if n <= 1:
        return 1
    # Recursive step (unchanged)
    return count_stairs(n - 1) + count_stairs(n - 2)

# Test with n=35 (runs instantly now!)
print(count_stairs(35))  # Output: 14930352
# Even try n=100—no performance issues at all
print(count_stairs(100)) # Output: 573147844013817084101

How the Closure Works (The Scope Magic) #

Let’s unpack what’s happening here:

1. When you decorate count_stairs with @memoize, Python runs memoize(count_stairs) once. This creates a memo dictionary that’s tied exclusively to count_stairs.
2. The memoize function returns the wrapper function, which replaces the original count_stairs in your code.
3. Every time you call count_stairs(n), you’re actually calling wrapper(n). Thanks to closure, wrapper retains access to the original memo dictionary even after memoize has finished running.
4. Recursive calls to count_stairs hit the same wrapper, so they reuse cached values instead of recomputing them—eliminating the exponential time complexity of the original code.

Making It More Flexible (For Multiple/Keyword Arguments) #

The above works for single positional arguments, but we can expand it to handle any function by creating a hashable key from all input parameters:

def memoize(func):
    memo = {}
    
    def wrapper(*args, **kwargs):
        # Create a unique key from positional and keyword arguments
        key = (args, frozenset(kwargs.items()))
        if key not in memo:
            memo[key] = func(*args, **kwargs)
        return memo[key]
    
    return wrapper

Now this decorator works for functions with any number of parameters, like:

@memoize
def calculate_area(length, width):
    return length * width

print(calculate_area(5, 3))  # Caches the result
print(calculate_area(width=3, length=5))  # Reuses the cached value

Why This Is Better Than Global Variables #

• Encapsulation: Each decorated function gets its own cache—no cross-contamination between different functions.
• Cleanliness: Your core logic (like count_stairs) stays focused on solving the problem, not managing cache state.
• Maintainability: You can modify or replace the memoize decorator later without touching the functions it decorates.

This is essentially the core pattern that functools.lru_cache uses (with extra features like cache size limits), but building it yourself gives you a concrete understanding of the underlying architecture.

内容的提问来源于stack exchange，提问作者Yevhen Ivashchenko