Pythonで累積和・累積積(itertools.accumulate)
Pythonで累積和や累積積を生成するには、標準ライブラリitertoolsモジュールのaccumulate()
関数を使う。累積和・累積積に限らず、任意の関数を累積的に適用できる。
累積和・累積積はNumPyやpandasの関数・メソッドでも生成可能。以下の記事を参照。
itertools.accumulate()の基本的な使い方(累積和を生成)
itertools.accumulate()
の引数にはリストやタプルなどのイテラブルオブジェクトを指定する。イテレータが返されるので、そのままprint()
などで出力しても結果は表示されない。
import itertools
l = [1, 2, 3, 4, 5, 6]
print(itertools.accumulate(l))
# <itertools.accumulate object at 0x1080e1a80>
print(type(itertools.accumulate(l)))
# <class 'itertools.accumulate'>
for文で使うと結果を順に取得できる。
for i in itertools.accumulate(l):
print(i)
# 1
# 3
# 6
# 10
# 15
# 21
結果をリストで取得したい場合はlist()
を使えばよい。タプルにしたい場合はtuple()
。
print(list(itertools.accumulate(l)))
# [1, 3, 6, 10, 15, 21]
先頭に0
の要素を追加したい場合は[0]
と連結する。Python3.8以降は後述の引数initial
を使う方法もある。
print([0] + list(itertools.accumulate(l)))
# [0, 1, 3, 6, 10, 15, 21]
逆順に累積和を生成したい場合は、元のリストをreversed()
で逆順にすればよい。
print(list(itertools.accumulate(reversed(l))))
# [6, 11, 15, 18, 20, 21]
適用する関数を指定(累積積などを生成): 引数func
要素に順に適用していく関数を引数func
で指定できる。デフォルトは加算で、これまでの例のように累積和が生成される。
例えば、引数func
に乗算の関数であるoperator.mul()
を指定すると累積積になる。関数などの呼び出し可能(callable)オブジェクトを引数に指定する場合は括弧()
をつけてはいけないので注意。
import itertools
import operator
print(operator.mul(2, 3))
# 6
l = [1, 2, 3, 4, 5, 6]
print(list(itertools.accumulate(l, func=operator.mul)))
# [1, 2, 6, 24, 120, 720]
operatorは+
や*
などの演算子に相当する関数を提供するモジュール。
加算(累積和)、乗算(累積積)だけでなく、減算や除算を累積的に適用することも可能。
print(list(itertools.accumulate(l, func=operator.sub)))
# [1, -1, -4, -8, -13, -19]
print(list(itertools.accumulate(l, func=operator.truediv)))
# [1, 0.5, 0.16666666666666666, 0.041666666666666664, 0.008333333333333333, 0.001388888888888889]
2つの引数を受け取る関数であれば、ほかの関数でもよい。
print(list(itertools.accumulate([1, 3, 2, 6, 5, 4], func=max)))
# [1, 3, 3, 6, 6, 6]
ラムダ式でもよい。任意の処理を適用できる。
print(list(itertools.accumulate(l, func=lambda x, y: x * y)))
# [1, 2, 6, 24, 120, 720]
print(list(itertools.accumulate(l, func=lambda x, y: int(str(x) + str(y)))))
# [1, 12, 123, 1234, 12345, 123456]
前者はoperator.mul()
と等価。後者は文字列に変換して+
で連結したあと、整数int
に戻している。
初期値を指定: 引数initial
Python3.8から引数initial
で初期値を設定できるようになった。
initial
に値を指定した場合、その値が最初の要素として扱われ、結果の要素数が一つ増える。デフォルトはinitial=None
。
import itertools
l = [1, 2, 3, 4, 5, 6]
print(list(itertools.accumulate(l)))
# [1, 3, 6, 10, 15, 21]
print(list(itertools.accumulate(l, initial=0)))
# [0, 1, 3, 6, 10, 15, 21]
print(list(itertools.accumulate(l, initial=100)))
# [100, 101, 103, 106, 110, 115, 121]