从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。它接受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]