NumPyで多項式計算

Python

Problem 389

これは確率計算をする問題なので、いつものように母関数の計算をすることになります。しかし、多項式をリストとして計算すると、Pythonは非常に遅いんですよね。そんなとき、NumPyを使うと劇的に速くなります。

ダウンロードはこのあたりから。
http://sourceforge.net/projects/numpy/files/

リファレンス
http://docs.scipy.org/doc/numpy/reference/generated/numpy.poly1d.html

まずは簡単な例から。

import numpy

f = numpy.poly1d([3, 2, 1])
print f

これを実行するとこう表示されます。

   2
3 x + 2 x + 1

x2 + 2x + 3ではないので注意です。

print f.c

とすると、

[3 2 1]

と表示されます。これは係数をNumPyのarrayにします。

print list(reversed(f.c))

とすれば、リストになります。

[1, 2, 3]

演算もふつうにできます。

print (f * f + 2 * f ).c

[ 9 12 16  8  3]


さて、速度を調べてみましょう。

(0.01 + 0.01x + ... + 0.01x99)100

をバイナリ法を使わずに計算します。

import numpy
import operator
import time

def mul_p(f, g):
    h = [ 0 ] * (len(f) + len(g) - 1)
    for k in xrange(len(f)):
        for l in xrange(len(g)):
            h[k+l] += f[k] * g[l]
    return h

t0 = time.clock()
N = 100
M = 100
f = [ 1.0 / N ] * N
g = reduce(mul_p, [ f ] * (M - 1), f)
print g[N*M/2]
print time.clock() - t0

t0 = time.clock()
f = numpy.poly1d([ 1.0 / N ] * N)
g = reduce(operator.mul, [ f ] * (M - 1), f)
print g[N*M/2]
print time.clock() - t0

NumPyを使わなかったとき13.6秒、使ったとき0.06秒でした。200倍以上速くなりましたね。

さて、これを使ってProblem 389を解くと、ナイーブなコードで7分が6秒、最後に工夫を入れたコードで0.9秒が0.06秒になりました。