如何生成符合规则的递增分组列？Pandas实现ExpectedGroup方法

需要为数据集生成ExpectedGroup列，规则如下：

• 当case_id为True时，该列值保持同一组的编号不变
• 当case_id为False时，该列值设为NaN（仅在连续的True块中递增组号）

已创建布尔掩码case_id，但使用np.where未得到预期结果，相关代码及数据集如下：

import pandas as pd
import numpy as np

df = pd.DataFrame([
            ['A', 'P', 'O', 2, np.nan],
            ['A', 'O', 'O', 5, 1],
            ['A', 'O', 'O', 10, 1],
            ['A', 'O', 'P', 4, np.nan],
            ['A', 'P', 'P', 300, np.nan],
            ['A', 'P', 'O', 2, np.nan],
            ['A', 'O', 'O', 5, 2],
            ['A', 'O', 'O', 10, 2],
            ['A', 'O', 'P', 4, np.nan],
            ['A', 'P', 'P', 300, np.nan],
            ['B', 'P', 'O', 2, np.nan],
            ['B', 'O', 'O', 5, 3],
            ['B', 'O', 'O', 10, 3],
            ['B', 'O', 'P', 4, np.nan],
            ['B', 'P', 'P', 300, np.nan],
            ],
        columns = ['ID', 'FromState', 'ToState', 'Hours', 'ExpectedGroup'])

# 创建布尔掩码
df['case_id'] = (df.FromState == 'O') & (df.ToState == 'O')

# 输出case_id结果：
# 0     False
# 1      True
# 2      True
# 3     False
# 4     False
# 5     False
# 6      True
# 7      True
# 8     False
# 9     False
# 10    False
# 11     True
# 12     True
# 13    False
# 14    False
# Name: case_id, dtype: bool

# 尝试的方法（未得到预期结果）
np.where(df['case_id'] != False, df['case_id'].cumsum(), np.nan)

正确实现方式 #

直接对case_id累加会把所有True行单独计数，无法形成连续相同的组号。正确思路是先标记每个连续True块的起始位置，再累加起始标记得到组号，最后仅在True行保留组号：

# 1. 标记每个连续True块的起始行
# 当前行case_id为True，且上一行不是True（首行特殊处理）
df['group_start'] = df['case_id'] & ~df['case_id'].shift(1, fill_value=False)

# 2. 累加起始标记，得到每组的编号
df['group_num'] = df['group_start'].cumsum()

# 3. 生成最终的ExpectedGroup列：True行保留组号，False行设为NaN
df['ExpectedGroup'] = np.where(df['case_id'], df['group_num'], np.nan)

# 查看结果
print(df[['ID', 'FromState', 'ToState', 'case_id', 'ExpectedGroup']])

运行后得到的ExpectedGroup列会完全匹配预期值：

• 行1、2为1，行6、7为2，行11、12为3，其余行均为NaN。

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