[memoization]Memoization decorator #python #optimization

memoization.py

class Memoize:
    def __init__(self, func): 
        self.func = func
        self.cache = {}
    def __call__(self, arg):
        if arg not in self.cache: 
            self.cache[arg] = self.func(arg)
        return self.cache[arg]


      
@Memoize
def partition_func(n):
    if n == 0: return 1
    result = 0; k = 1; sign = 1; 
    while True:
        pent = (3*k**2 - k) // 2        
        if pent > n: break
        result += sign * partition_func(n - pent) 
        pent += k
        if pent > n: break
        result += sign * partition_func(n - pent) 
        k += 1; sign = -sign; 
    return result