记忆函数装饰器-避免函数重复计算的利器

注意,此装饰器的设计前提是:对于相同参数,函数一定返回相同的值.如果某个函数,相同参数可能返回不同值,或者修改了一些外部数据,那么最好不要用此装饰器.

有时候对于一些函数,我们需要用相同参数多次调用,且不方便将其预先储存在一个变量中.

例如这种,假设有函数f,g,h,m:

list=[[f(i),g(f(i)),h(g(f(i))),m(f(i))] for i in range(100)]

那么,在计算g(f(i)),要先重复计算一次f(i),再计算g.

h(g(f(i)))同理.如果这些函数的计算过程很复杂,那么就浪费了大量资源.

再例如很经典的,非尾递归形式的斐波拉契函数:

def fibonacci(n):
    print "call fibonacci"
    if n in (0, 1):
        return n
    return fibonacci(n-1) + fibonacci(n-2)

你想要计算fibonacci(100),你会发现非常非常慢.因为它最终等于计算一大堆fibonacci(1)和fibonacci(0)之和.简直是重复计算的典范了.

这还是第一层重复计算.如果非常不幸,你还要重复调用fibonacci(100),那么简直无法想象后果.

对于这种浪费,可以让函数自己记忆每次被调用的参数和结果信息.

为了方便应用,我们可以设计一个函数装饰器,凡是经过它装饰的函数,在被调用时都会首先检查传入的参数是不是以前有过的,如果是,那么直接返回结果;如果不是,那么就进行计算,保存此次参数和函数值,再返回结果.

def memory(function):
    cache = {}
    def memofunc(*nkw,**kw):
        key=str(nkw)+str(kw)
        if key not in cache:            
            cache[key] = function(*nkw,**kw)
        return cache[key]
    return memofunc
    

@memory
def fibonacci(n):
    print "call fibonacci"
    if n in (0, 1):
        return n
    return fibonacci(n-1) + fibonacci(n-2)


@memory
def func(a,b,*nkw,**kw):
    print 'call func'
    return a+b+sum(nkw)+sum(kw.values())

print '-'*20
print fibonacci(10)
print '-'*20
print fibonacci(10)
print '-'*20
print func(1,2,3,4,5,x=6,y=7)
print '-'*20
print func(1,2,3,4,5,x=6,y=7)

最终结果:

>>> 
--------------------
call fibonacci
call fibonacci
call fibonacci
call fibonacci
call fibonacci
call fibonacci
call fibonacci
call fibonacci
call fibonacci
call fibonacci
call fibonacci
55
--------------------
55
--------------------
call func
28
--------------------
28

测试发现,上面的装饰器可以让fib在1秒内算出fib(332).更大的数就会报错了.

大家可以比较一下去掉那个装饰器之后,Python要多久才算得出fib(332).

此外,对于一些只接受1个参数的函数来说,还可以有个简单的装饰器版本,性能会更高一些:

def memory(function):
    cache = {}
    def memofunc(x):
        if x not in cache:
            cache[x] = function(x)
        return cache[x]
    return memofunc

用这个装饰器,可以计算到fib(498)