python - python 中是否有类似于 `numpy.cumsum` 或 `numpy.cumprod` 的函数，如`cumreduce`？

Question

当我需要列表中元素的总和或乘积时，例如

>>> foo = [10, 5, 3, 4]

我可以使用numpy的sumor函数prod

>>> import numpy as np

>>> np.sum(foo)
22

>>> np.prod(foo)
600

同样，当我需要累积总和或乘积时，我可以使用np.cumsum或np.cumprod

>>> np.cumsum(foo)
array([10, 15, 18, 22])

>>> np.cumprod(foo)
array([ 10,  50, 150, 600])

有没有办法获得任意reduce操作的累积结果？

---

例如，如果我有一个像

def my_fn(a, b):
    return a + b**2

我可以functools.reduce用来获得

>>> from functools import reduce
>>> reduce(my_fn, foo)
60

我正在寻找的是一个功能，它会给我以下

>>> cumreduce(my_fn, foo)
[10, 35, 44, 60]

也就是说，结果的每个元素都等于reduce(my_fn, foo[:i])。

---

当然，我可以以天真的方式做到这一点

>>> def cumreduce(fn, seq):
...     return [reduce(fn, seq[:i]) for i in range(1, len(seq)+1)]

>>> cumreduce(my_fn, foo)
[10, 35, 44, 60]

理想情况下，我正在寻找具有相同或相似功能的内置函数。

标签： pythonnumpy

Answer 1

在 Python 中，您正在寻找的是itertools.accumulate：

import itertools

[*itertools.accumulate(foo,my_fn)]
# [10, 35, 44, 60]

Numpy ufuncs 通常有一个累积方法。例子：

np.bitwise_xor.accumulate(foo)
array([10, 15, 12,  8])

np.add.accumulate(foo)
array([10, 15, 18, 22])
# cf. cumsum:
np.cumsum(foo)
array([10, 15, 18, 22]