Python memoize decorator

偶然看 Python 代码,看到这么一个叫memoize decorator 的东西,你会发现它非常有用。

看下面的例子,我们最常见的 fibonacci 数列的递归实现。

~% python
>>> def fib(n):
...     if n in (0, 1):
...             return n
...     else:
...             return fib(n-1)+fib(n-2)
...
>>> import timeit
>>> timeit.timeit("fib(5)", "from __main__ import fib")
3.946546792984009
>>> timeit.timeit("fib(7)", "from __main__ import fib")
10.944302082061768

很慢,不是么?你当然可以把它的实现改成迭代,但是,如果前提是我们不对fib()函数本身做任何代码修改的话,我们还有别的方法去改进它。

我们注意到,fib(n-1) 和 fib(n-2) 实际上重复计算了很多结果,根本没必要,一个自然而然的想法就是把这些结果缓存起来,于是就有了下面的:

>>> def memoize(f):
...     memo = {}
...     def wrapper(x):
...         if x not in memo:
...             memo[x] = f(x)
...         return memo[x]
...     return wrapper
...
>>>
>>> def fib(n):
...     if n in (0, 1):
...         return n
...     else:
...         return fib(n-1) + fib(n-2)
...
>>> fib = memoize(fib)
>>> import timeit
>>> timeit.timeit("fib(7)", "from __main__ import fib")
0.35535502433776855

还不够,Python 还有自己的解决办法,那就是用 decorator

[python]
def memoize(function):
memo = {}
def wrapper(args):
if args in memo:
return memo[args]
else:
rv = function(
args)
memo[args] = rv
return rv
return wrapper

@memoize
def fib(n):
if n in (0, 1):
return n
return fib(n-1) + fib(n-2)
[/python]

为此,在 Python3 中还专门引入了一个 lru_cache

这个东西的牛逼之处在于,有了它你可以放心地去写递归的代码,不用为了效率去把它改成迭代形式,你也知道,有些时候递归代码比迭代代码更容易理解。这样可读性有了,效率也提升了,不用修改函数的实现,直接让它从次方级降到了线性级。