Unique Binary Trees @ LeetCode

ubt.py

"""
class Solution {
public:
    int numTrees(int n) {
        if (n == 1)
            return 1;
        
        int res = 0;
        // for each of the middle nodes, the number is the product of its left and right children
        for (int i = 2; i <= n - 1; ++i)
            res += numTrees(i - 1) * numTrees(n - i);
        
        // the two 'degraded' trees without left/right child
        res += 2 * numTrees(n - 1);
        return res;
    }
};
"""

# def memodict(f):
#    """ Memorization dict for a single argument function """
#    class memodict(dict):
#        def __missing__(self, key):
#            res = self[key] = f(key)
#            return res
#    
#    return memodict().__getitem__

from functools import partial

class memo(object):
    def __init__(self, func):
        self.func = func
    
    def __get__(self, obj, objtype=None):
        if obj is None:
            return self.func
        return partial(self, obj)
    
    def __call__(self, *args, **kw):
        obj = args[0]
        try:
            cache = obj.__cache
        except AttributeError:
            cache = obj.__cache = {}
        key = (self.func, args[1:], frozenset(kw.items()))
        try:
            res = cache[key]
        except KeyError:
            res = cache[key] = self.func(*args, **kw)
        return res
            

class Solution:
    @memo
    def numTrees(self, n):
        if n == 1:
            return 1
        
        res = 0
        for i in range(2, n):
            res += self.numTrees(i - 1) * self.numTrees(n - i)
        res += 2 * self.numTrees(n - 1)
        return res