条件运行计数熊猫

时间:2017-09-08 17:32:01

标签: python pandas group-by

我试图根据两个条件在pandas中创建条件运行总和。

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']

数据如下:

   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

然后我想看看ID之前有多长时间为B,我的尝试:

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

这给了我:

    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

理想输出如下:

   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

正如你所看到的那样,ID更改我希望time_on_b重置,所以它给我的值为1而不是3.

2 个答案:

答案 0 :(得分:4)

似乎您需要按ID分组,然后使用cumsum来计算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

答案 1 :(得分:2)

你也可以使用transform之类的

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