有没有 reset 的 cumsum 函数

[admirez]

比如［-1，-1，-1，1，1，1，1，-1，-1 ］。

能变成 -1，-2，-3，1，2，3，4，-1，-2


最好是高效率的

[necomancer]

a = np.array([-1,-1,-1,1,1,1,1,-1,-1])
bounds = np.r_[0, np.argwhere(np.diff(a>0, prepend=0)).ravel(), a.shape[0]]
np.concatenate([np.cumsum(a[i:j]) for i, j in zip(bounds[:-1], bounds[1:])])

[necomancer]

可能快一点，尤其是 1， -1 交错很厉害的时候
import pandas as pd
b = pd.Series(np.insert(a, np.argwhere(np.diff(a>0, prepend=0)).ravel(), 0))
cum = b.cumsum()
cum = (cum -cum.where(b==0).ffill().fillna(0)).values
cum = cum[cum!=0]

[admirez]

@necomancer 非常感谢，我试试看