根据条件计算数组中的累积总和

问题描述

在我使用的另一种语言中，有一个名为 SumSince( condition,array ) 的函数。所以例如我有一个 DataArray 和一个 ConditionArray 看起来像

数据数组	条件数组	结果数组
1	1	1
2	0	3
3	0	6
4	1	4
5	0	9
6	0	15

然后 ResultArray = SumSince( ConditionArray,DataArray ) 将给出第 3 列中显示的数组

所以它从条件开始累积。当它满足一个新条件时，它将从该条件开始从头开始累积。

Python 中是否有类似的功能。我找不到。我知道我可以使用循环，但我很确定某处有一个函数可以做到这一点

谢谢

解决方法

我暂时只是编写了自己的函数。在 Pandas 中，你有 Cumsum。我只是没有找到如何将它用作 SumSince。我在下面写了一个函数。似乎工作。可能是丑陋的代码，但刚刚开始使用 Python。如果有人知道已经存在的功能，请告诉我，谢谢

def SumSince(condition,array):
    length = len(array)
    outputArray = [0] * length
    i = 0
    csum = 0
    while i < length:
        if (condition[i] == 0):
            csum = array[i] + csum
            outputArray[i] = csum
        elif (condition[i] == 1):
            csum = array[i]
            outputArray[i] = csum
        i += 1
    return outputArray

array = [1,2,3,4,5,6]
condition = [1,1,0]
tst = SumSince(condition,array)

tst
[1,6,9,15]