リストから和のリストを得る（累積和）

はじめに

数値計算をしていると、積分のために総和を取ることはよくある。たとえば、物理でメッシュごとの電場のデータからポテンシャルを求めるには、総和を取ってメッシュ幅をかける。

リストAがあるとして、総和を取るだけならsum(A)とすればよい。だが「リストAのi番目までの和からなるリスト」を得るにはどのようにすればよいだろうか。

今回はいくつか思いつく方法を試して実行時間を計測してみた。

1.for文

多分これが一番すぐに思いつくと思う。

res[0] = mydata[0]
for i in range(len(mydata)-1):
    res[i+1] = res[i] + mydata[i+1]

2.リスト内包表記

ちょっと調子に乗ったPythonの哲学を考えた人は「for文は使いたくない」と思うだろう。for文の回避に使う典型的なものといえばリスト内包表記だ。

res = [sum(mydata[:i+1]) for i in range(len(mydata))]

ただし、少し考えてみればすぐわかるが、これは一番悪手になる。なぜなら絶対的な計算量が増えるからだ。i番目の要素を計算するたびにそこまでの足し算を行うので$\frac{1}{2}n(n+1)$回の足し算を行うことになってしまう。for文なら$n$回だ。mapとlambdaで処理するのも同じだ。

3.再帰

for文を避ける方法はもう一つ、再帰を使うという方法がある。

import sys
sys.setrecursionlimit(len(mydata)+10)

def step(myres,mymydata,i):
    if i == len(mymydata)-1:
        pass
    else:
        myres[i+1] = myres[i] + mymydata[i+1] 
        step(myres,mymydata,i+1)

step(res,mydata,0)

Pythonには無限ループに陥らないようにするため再帰の上限回数が定められており、デフォルトでは100回くらいで止まってしまう。そのためsys.setrecursionlimit()で上限回数を設定している。

確かにfor文は回避できたが、明らかにコードが複雑になる。

4.NumPy.cumsum()

お察しかもしれないが、NumPyにはcumsumという関数がある。調べてから知ったのだが扱っているような演算を累積和というらしい。

import numpy
res = numpy.cumsum(mydata)

計測結果

以上を踏まえ、実行時間を計測してみた。

結果：NumPy.cumsum()の圧勝。Numpy.cumsumを使いましょう。

計測に使ったコード

下記のgistまで。
https://gist.github.com/Bluepost59/85ce18cdfde26c31018ce018488ae07d