CODEWITHSUNDEEP
pythonpandasgroup-by
Adam Warner
Adam Warner

Reputation: 1354

Conditional Running Count Pandas

I am attempting to create a conditional running sum in pandas based on two conditions. 

import pandas as pd
ID = [1,1,1,2,2,3,4]
after = ['A','B','B','A','A','B','A']
before = ['A','B','B','A','A','B','A']
df = pd.DataFrame([ID, before,after]).T
df.columns = ['ID','before','after']

The data looks like: 

   ID before after
0  1      A     A
1  1      B     B
2  1      B     B
3  2      A     A
4  2      A     A
5  3      B     B
6  4      A     A

I then want to look at how long an ID has had the value of before as B, my attempt:

df['time_on_b'] = (df.groupby('before')['ID'].cumcount()+1).where(df['before']=='B',0)

This gives me: 

    ID before after  time_on_b
0  1      A     A          0
1  1      B     B          1
2  1      B     B          2
3  2      A     A          0
4  2      A     A          0
5  3      B     B          3
6  4      A     A          0

The ideal output is as followed: 

   ID before after  time_on_b
0  1      A     A          0
1  1      B     B          1
2  1      B     B          2
3  2      A     A          0
4  2      A     A          0
5  3      B     B          1
6  4      A     A          0

As you can see as the ID changes I want time_on_b to reset so it gives me the value of 1 and not 3.

Upvotes: 3

Views: 1785

Answers (2)

piroot
piroot

Reputation: 772

You could also use transform like

df.groupby('ID').before.transform(lambda x: x.eq('B').cumsum())

0    0
1    1
2    2
3    0
4    0
5    1
6    0
Name: before, dtype: int32

df.assign(time_on_b=df.groupby('ID').before.transform(lambda x: x.eq('B').cumsum()))

   ID before after  time_on_b
0   1      A     A          0
1   1      B     B          1
2   1      B     B          2
3   2      A     A          0
4   2      A     A          0
5   3      B     B          1
6   4      A     A          0

Upvotes: 3

akuiper
akuiper

Reputation: 215047

Seems you need to group by ID, then use cumsum to count the occurrences of B:

cond = df.before == 'B'
df['time_on_b'] = cond.groupby(df.ID).cumsum().where(cond, 0).astype(int)
df
#  ID   before  after   time_on_b
#0  1        A      A   0
#1  1        B      B   1
#2  1        B      B   2
#3  2        A      A   0
#4  2        A      A   0
#5  3        B      B   1
#6  4        A      A   0

Upvotes: 7

Related Questions