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

Question

我有一个数字序列，例如 [0,0,0,0,1,1,1,0,0,1,1,0]。我想根据最后一个非零值计算数字总和。即，一旦出现零条目，cumsum 将重置为零。

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, elem: acc + elem if elem else 0)
print(list(result))
# [0, 0, 0, 0, 1, 2, 3, 0, 0, 1, 2, 0]