python - 条件运行计数 Pandas

Question

我正在尝试根据两个条件在 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。

Answer 1

看来您需要按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