每次有0时都会重置的Python pandas cumsum

Question

我有一个0和1的矩阵，想对每一个列做一个累加，每当观察到一个零时就重置为0。例如，如果我们具有以下内容：

df = pd.DataFrame([[0,1],[1,1],[0,1],[1,0],[1,1],[0,1]],columns = ['a','b'])
print(df)
   a  b
0  0  1
1  1  1
2  0  1
3  1  0
4  1  1
5  0  1

我想要的结果是：

print(df)
   a  b
0  0  1
1  1  2
2  0  3
3  1  0
4  2  1
5  0  2

但是，当我尝试时df.cumsum() * df，我能够正确识别0个元素，但计数器不会重置：

print(df.cumsum() * df)
   a  b
0  0  1
1  1  2
2  0  3
3  2  0
4  3  4
5  0  5

Answer 1

您可以使用：

a = df != 0
df1 = a.cumsum()-a.cumsum().where(~a).ffill().fillna(0).astype(int)
print (df1)
   a  b
0  0  1
1  1  2
2  0  3
3  1  0
4  2  1
5  0  2