CODEWITHSUNDEEP
pythonpandas
Andrew Vitek
Andrew Vitek

Reputation: 37

Using previous row value while creating a new column

I have a df in python that looks something like this:

'A'
 0 
 1
 0
 0
 1
 1
 1
 1
 0

I want to create another column that adds cumulative 1's from column A, and starts over if the value in column A becomes 0 again. So desired output:

'A' 'B'
 0   0
 1   1
 0   0
 0   0
 1   1
 1   2
 1   3
 1   4
 0   0

This is what I am trying, but it's just replicating column A:

df.B[df.A ==0] = 0
df.B[df.A !=0] = df.A + df.B.shift(1)

Upvotes: 1

Views: 43

Answers (2)

Space Impact
Space Impact

Reputation: 13255

Use shift with ne and groupby.cumsum:

df['B'] = df.groupby(df['A'].shift().ne(df['A']).cumsum())['A'].cumsum()

print(df)
   A  B
0  0  0
1  1  1
2  0  0
3  0  0
4  1  1
5  1  2
6  1  3
7  1  4
8  0  0

Upvotes: 1

BENY
BENY

Reputation: 323396

Let us do cumsum with groupby cumcount

df['B']=(df.groupby(df.A.eq(0).cumsum()).cumcount()).where(df.A==1,0)
Out[81]: 
0    0
1    1
2    0
3    0
4    1
5    2
6    3
7    4
8    0
dtype: int64

Upvotes: 2

Related Questions