计算Python中最后一个非零条目的累积和

Question

我有一个数字系列，例如 [0,0,0,0,1,1,1,0,0,1,1,0]。我想计算最后一个非零值的数字总和。即一旦发生零输入，累积和将重置为零。

input: [0,0,0,0,1,1,1,0,0,1,1,0]
output:[0,0,0,0,1,2,3,0,0,1,2,0] 

有内置的python函数可以实现这个功能吗？或者更好的方法来计算它而不循环？

Answer 1

您可以使用

itertools.accumulate

来完成。除了传递一个可迭代对象作为第一个参数之外，它还接受一个可选的第二个参数，该参数应该是一个双参数函数，其中第一个参数是累积结果，第二个参数是可迭代对象的当前值。您可以传递一个相当简单的 lambda 作为可选的第二个参数来计算运行总计，除非当前值为零。

from itertools import accumulate

nums = [0,0,0,0,1,1,1,0,0,1,1,0]

result = accumulate(nums, lambda acc, n: acc + n if n else 0)
print(list(result))
# [0, 0, 0, 0, 1, 2, 3, 0, 0, 1, 2, 0]

Answer 2

我们可以在 numpy 中通过两次

np.cumsum(..)

来完成此操作。首先我们计算数组的

cumsum

：

a = np.array([0,0,0,0,1,1,1,0,0,1,1,0])
c = np.cumsum(a)

这给了我们：

>>> c
array([0, 0, 0, 0, 1, 2, 3, 3, 3, 4, 5, 5])

接下来，我们对值为

a

的元素过滤

0

，并按元素计算该元素与其前一个元素之间的差异：

corr = np.diff(np.hstack(((0,), c[a == 0])))

那么这就是我们需要对这些元素应用的修正：

>>> corr
array([0, 0, 0, 0, 3, 0, 2])

然后我们可以复制

a

（或就地执行此操作），并减去校正：

a2 = a.copy()
a2[a == 0] -= corr

这给了我们：

>>> a2
array([ 0,  0,  0,  0,  1,  1,  1, -3,  0,  1,  1, -2])

现在我们可以计算

a2

的累积和，对于

0

，它将重置为

0

，因为校正会跟踪其间的增量：

>>> a2.cumsum()
array([0, 0, 0, 0, 1, 2, 3, 0, 0, 1, 2, 0])

或作为函数：

import numpy as np

def cumsumreset(iterable, reset=0):
    a = np.array(iterable)
    c = a.cumsum()
    a2 = a.copy()
    filter = a == reset
    a2[filter] -= np.diff(np.hstack(((0,), c[filter])))
    return a2.cumsum()

这给了我们：

>>> cumsumreset([0,0,0,0,1,1,1,0,0,1,1,0])
array([0, 0, 0, 0, 1, 2, 3, 0, 0, 1, 2, 0])