如何向熊猫数据框添加条件计数器

Question

我有一个如下图所示的数据框：

id  date       notify
3   04/09/2019  no
3   30/10/2019  yes
3   03/05/2020  no
3   05/09/2020  no
3   31/10/2020  yes
3   03/11/2020  no
5   03/09/2019  no
5   27/10/2019  yes
5   02/05/2020  no

我想在每次“通知”为“是”时创建一个计数器组编号。 然后我想将相同的数字应用于“通知”将始终为“否”的下一行。

应该是这样的：

id  date       notify time_group
3   04/09/2019  no       
3   30/10/2019  yes       1
3   03/05/2020  no        1
3   05/09/2020  no
3   31/10/2020  yes       2
3   03/11/2020  no        2
5   03/09/2019  no
5   27/10/2019  yes       3
5   02/05/2020  no        3

目前我已经尝试过了，但没有成功：

i = 0
df['time_grp'] = np.nan
for row in df.iterrows():
    if row['notify'] == 'yes':
        row['time_group'] = i
        i += 1

我想知道是否有更适合熊猫的方式来完成此任务？也许利用 cumcount()？我知道我可以使用 shift(-1) 将相同的计数器组号应用于下一行...

Answer 1

试试：

# mark the `yes` rows
s = df['notify'].eq('yes')


# s.cumsum() enumerate the blocks
# maybe `s.groupby(df['id']).cumsum() if enumeration within id
df['time_group'] = s.cumsum().where(               # use `where` to keep      
     s |                                           # the `yes` rows
     s.groupby(df['id']).shift(fill_value=False)   # and those after
)