如何在Pandas按组从第二个'outcome2'开始累计计数？

问题描述 #

需要在Pandas DataFrame中新增一列，满足以下要求：

• 按datetime顺序统计value列中'outcome2'的出现次数
• 从每组第二次观测到'outcome2'时开始累计计数
• 按ID分组统计

示例数据与初始代码 #

import pandas as pd
from io import StringIO

txt= """
ID,datetime,value
A,12/10/2022 10:00:00,outcome1
A,12/10/2022 11:15:10,outcome2
A,14/10/2022 15:30:30,outcome1
B,11/10/2022 11:30:22,outcome1
B,15/10/2022 22:44:11,outcome2
B,15/10/2022 23:30:22,outcome3
B,15/10/2022 23:31:11,outcome2
"""

df = pd.read_csv(StringIO(txt),
                 parse_dates=[1],
                 dayfirst=True)
                 .assign(id_index= lambda x_df: x_df
                 .groupby('ID', sort=False).ngroup())
                 .set_index("id_index")
                 .rename_axis(index=None)

当前尝试与问题 #

初始尝试生成了中间列value_test和value_cumsum，得到如下结果：

ID            datetime     value value_test   value_cumsum
0  A 2022-10-12 10:00:00  outcome1      False              0
0  A 2022-10-12 11:15:10  outcome2       True              1
0  A 2022-10-14 15:30:30  outcome1      False              1
1  B 2022-10-11 11:30:22  outcome1      False              0
1  B 2022-10-15 22:44:11  outcome2       True              1
1  B 2022-10-15 23:30:22  outcome3      False              1
1  B 2022-10-15 23:31:11  outcome2       True              2

后续尝试用嵌套lambda生成目标列，虽能运行但代码冗余：

df = df.assign(value_test = lambda df: df['value']=='outcome2',
               cumsum = lambda df: df.groupby('ID', sort=False)
                        ['value_test'].cumsum(),
               outcome2 = lambda df:df['cumsum'].apply(
                          lambda cumsum: 0 if cumsum == 1
                                 else (0 if cumsum == 0
                                 else (cumsum-1 if cumsum > 1
                                 else 'NaN'))))

期望结果 #

希望得到无需中间变量的简洁实现，最终DataFrame如下：

ID            datetime     value outcome2
0  A 2022-10-12 10:00:00  outcome1        0
0  A 2022-10-12 11:15:10  outcome2        0
0  A 2022-10-14 15:30:30  outcome1        0
1  B 2022-10-11 11:30:22  outcome1        0
1  B 2022-10-15 22:44:11  outcome2        0
1  B 2022-10-15 23:30:22  outcome3        0
1  B 2022-10-15 23:31:11  outcome2        1

最优实现方案 #

可以直接通过链式调用完成，无需任何中间变量，代码简洁且高效：

方法一（高效矢量运算） #

df['outcome2'] = (
    (df['value'] == 'outcome2')  # 标记是否为目标值
    .groupby(df['ID'], sort=False)  # 按ID分组
    .cumsum()  # 分组累计计数
    .sub(1)  # 累计值减1，让第一次出现的计数变为0
    .clip(lower=0)  # 将负数（非目标值或第一次出现前的行）裁剪为0
)

方法二（assign链式调用） #

如果希望保持assign的链式风格，可以这样写：

df = df.assign(
    outcome2=lambda x: (x['value'] == 'outcome2')
                       .groupby(x['ID'], sort=False)
                       .cumsum()
                       .sub(1)
                       .clip(lower=0)
)

实现逻辑说明 #

1. 标记目标值：df['value'] == 'outcome2'生成布尔序列，True代表当前行是目标值。
2. 分组累计：按ID分组后用cumsum()统计每组内目标值的出现次数（第一次出现为1，第二次为2，以此类推）。
3. 调整计数起点：用sub(1)将累计值减1，此时第一次出现的目标值计数变为0，第二次变为1，正好符合"从第二次开始累计"的需求。
4. 修正非目标值行：用clip(lower=0)将非目标值行的负数（累计和为0时减1得到-1）修正为0，保证所有非计数行的结果都是0。

最终运行结果 #

执行上述代码后，得到的DataFrame与期望结果完全一致：

ID            datetime     value  outcome2
0  A 2022-10-12 10:00:00  outcome1         0
0  A 2022-10-12 11:15:10  outcome2         0
0  A 2022-10-14 15:30:30  outcome1         0
1  B 2022-10-11 11:30:22  outcome1         0
1  B 2022-10-15 22:44:11  outcome2         0
1  B 2022-10-15 23:30:22  outcome3         0
1  B 2022-10-15 23:31:11  outcome2         1

内容的提问来源于stack exchange，提问作者Johan