优化分区函数

Question

这里有一段用Python写的代码：

# function for pentagonal numbers
def pent (n):     return int((0.5*n)*((3*n)-1))

# function for generalized pentagonal numbers
def gen_pent (n): return pent(int(((-1)**(n+1))*(round((n+1)/2))))

# array for storing partitions - first ten already stored
partitions = [1, 1, 2, 3, 5, 7, 11, 15, 22, 30, 42]

# function to generate partitions 
def partition (k):
 if (k < len(partitions)): return partitions[k]

 total, sign, i = 0, 1, 1

 while (k - gen_pent(i)) >= 0:
  sign   = (-1)**(int((i-1)/2))
  total += sign*(partition(k - gen_pent(i)))
  i     +=  1

 partitions.insert(k,total)
 return total

它使用这个方法来计算分区：

p(k) = p(k − 1) + p(k − 2) − p(k  − 5) − p(k − 7) + p(k − 12) + p(k  − 15) ...

(来源：维基百科)

不过，当处理大数字（超过p(10^3)）时，这段代码运行得非常慢。我想请教一下，是否可以优化我的代码，或者给我一些更快的算法或方法的建议。任何优化的建议都非常欢迎。

另外，先说明一下，这不是作业或Project Euler的题目，只是为了积累经验。

Answer 1

这里有一些评论。请注意，我并不是这方面的专家，但我也喜欢玩数学（还有Project Euler）。

我重新定义了五边形数字的函数，具体如下：

def pent_new(n):
    return (n*(3*n - 1))/2

def gen_pent_new(n):
    if n%2:
        i = (n + 1)/2
    else:
        i = -n/2
    return pent_new(i)

我写这些函数的时候，避免使用浮点数计算，因为在处理大数字时，使用浮点数会引入错误（比如你可以对比一下 n = 100000001 的结果）。

你可以使用 timeit 模块来测试一些小代码片段。当我测试所有的五边形函数（你的和我的）时，我的函数运行得更快。下面是一个测试你 gen_pent 函数的例子。

# Add this code to your script
t = Timer("for i in xrange(1, 100): gen_pent(i)", "from __main__ import gen_pent")
print t.timeit()

这是你 partition 函数的一个小修改。同样，使用 timeit 测试显示它比你的 partition 函数更快。不过，这可能是因为我对五边形数字函数做了一些改进。

def partition_new(n):
    try:
        return partitions_new[n]
    except IndexError:
        total, sign, i = 0, 1, 1
        k = gen_pent_new(i)
        while n - k >= 0:
            total += sign*partition_new(n - k)

            i += 1
            if i%2: sign *= -1
            k = gen_pent_new(i)

        partitions_new.insert(n, total)
        return total

在你的 partition 函数中，你在每次循环中都计算了一次通用的五边形数字。一次是为了在 while 条件中测试，另一次是为了更新 total。我把结果存储在一个变量里，这样每次循环只计算一次。

使用 cProfile 模块（适用于 python 版本 >= 2.5，否则使用 profile 模块），你可以看到你的代码大部分时间花在哪里。下面是基于你 partition 函数的一个例子。

import cProfile
import pstats

cProfile.run('partition(70)', 'partition.test')
p = pstats.Stats('partition.test')
p.sort_stats('name')
p.print_stats()

这在控制台窗口中产生了以下输出：

C:\Documents and Settings\ags97128\Desktop>test.py
Fri Jul 02 12:42:15 2010    partition.test

         4652 function calls (4101 primitive calls) in 0.015 CPU seconds

   Ordered by: function name

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
      552    0.001    0.000    0.001    0.000 {len}
        1    0.000    0.000    0.000    0.000 {method 'disable' of '_lsprof.Profiler' objects}
       60    0.000    0.000    0.000    0.000 {method 'insert' of 'list' objects}
        1    0.000    0.000    0.015    0.015 <string>:1(<module>)
     1162    0.002    0.000    0.002    0.000 {round}
     1162    0.006    0.000    0.009    0.000 C:\Documents and Settings\ags97128\Desktop\test.py:11(gen_pent)
    552/1    0.005    0.000    0.015    0.015 C:\Documents and Settings\ags97128\Desktop\test.py:26(partition)
     1162    0.002    0.000    0.002    0.000 C:\Documents and Settings\ags97128\Desktop\test.py:5(pent)

现在对我的 partition 函数进行性能分析：

Fri Jul 02 12:50:10 2010    partition.test

         1836 function calls (1285 primitive calls) in 0.006 CPU seconds

   Ordered by: function name

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.000    0.000    0.000    0.000 {method 'disable' of '_lsprof.Profiler' objects}
       60    0.000    0.000    0.000    0.000 {method 'insert' of 'list' objects}
        1    0.000    0.000    0.006    0.006 <string>:1(<module>)
      611    0.002    0.000    0.003    0.000 C:\Documents and Settings\ags97128\Desktop\test.py:14(gen_pent_new)
    552/1    0.003    0.000    0.006    0.006 C:\Documents and Settings\ags97128\Desktop\test.py:40(partition_new)
      611    0.001    0.000    0.001    0.000 C:\Documents and Settings\ags97128\Desktop\test.py:8(pent_new)

你可以看到我的结果中没有调用 len 和 round 这两个内置函数。而且我几乎将对五边形函数（gen_pent_new 和 pent_new）的调用次数减半了。

可能还有其他方法可以提高 Python 代码的速度。我建议你在这里搜索“python 优化”，看看能找到什么。

最后，你提到的维基百科页面底部的一个链接是关于快速计算分区数的 学术论文。快速浏览一下可以看到它包含了算法的伪代码。这些算法可能比你或我能写的任何东西都要快。