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进行操作。它接受2自变量函数，其中第一个自变量是累加结果，第二个自变量是可迭代对象中的当前元素，因此，除非当前元素为零，否则您可以使用相当简单的lambda来计算运行总额。

from itertools import accumulate

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

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

Answer 2

我们可以通过两次np.cumsum(..)以numpy的方式进行操作。首先，我们计算数组的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的累积总和，因为校正会跟踪之间的增量：

或作为功能：

>>> 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()

从python中的最后一个非零条目计算累积总和

2 个答案: