Python pandas cumsum with reset everytime there is a 0 [duplicate]

Question

This question already has an answer here:

• Cumsum reset at NaN 4 answers

I have a matrix with 0s and 1s, and want to do a cumsum on each column that resets to 0 whenever a zero is observed. For example, if we have the following:

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

The result I desire is:

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

However, when I try df.cumsum() * df, I am able to correctly identify the 0 elements, but the counter does not reset:

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

You can use:

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

Answer 2

Try this

df = pd.DataFrame([[0,1],[1,1],[0,1],[1,0],[1,1],[0,1]],columns = ['a','b'])
df['groupId1']=df.a.eq(0).cumsum()
df['groupId2']=df.b.eq(0).cumsum()
New=pd.DataFrame()
New['a']=df.groupby('groupId1').a.transform('cumsum')
New['b']=df.groupby('groupId2').b.transform('cumsum')

New
Out[1184]: 
   a  b
0  0  1
1  1  2
2  0  3
3  1  0
4  2  1
5  0  2

Answer 3

A slightly hacky way would be to identify the indices of the zeros and set the corresponding values to the negative of those indices before doing the cumsum:

import pandas as pd
df = pd.DataFrame([[0,1],[1,1],[0,1],[1,0],[1,1],[0,1]],columns = ['a','b'])
z = np.where(df['b']==0)
df['b'][z[0]] = -z[0]
df['b'] = np.cumsum(df['b'])
df

   a  b
0  0  1
1  1  2
2  0  3
3  1  0
4  1  1
5  0  2