如何在Pandas DataFrame中按列分组并添加分组标识列G

需求说明 #

现有如下Pandas DataFrame：

import pandas as pd

data = {'A': ['ABCD_1', 'ABCD_1', 'ABCD_1', 'ABCD_1', 'PQRS_2', 'PQRS_2', 'PQRS_2', 'PQRS_2', 'PQRS_2'], 
        'P':[1, 0, 0, 0, 0, 1, 0, 0, 0]}
df = pd.DataFrame(data)

需要按A列分组，在每个组内P=1行之后的第二个P=0行的G列填入组序号（第一个组填1，第二个组填2），其余行G列留空。

实现代码 #

import pandas as pd

data = {'A': ['ABCD_1', 'ABCD_1', 'ABCD_1', 'ABCD_1', 'PQRS_2', 'PQRS_2', 'PQRS_2', 'PQRS_2', 'PQRS_2'], 
        'P':[1, 0, 0, 0, 0, 1, 0, 0, 0]}
df = pd.DataFrame(data)

# 初始化G列为空字符串
df['G'] = ''

group_number = 1
# 按A列分组遍历
for group_name, group_data in df.groupby('A'):
    # 找到当前组内P=1的行索引
    p1_row_idx = group_data[group_data['P'] == 1].index[0]
    # 筛选P=1行之后的所有P=0的行，取第二个的索引
    target_row_idx = group_data.loc[p1_row_idx:][group_data['P'] == 0].index[1]
    # 为目标行赋值组编号
    df.loc[target_row_idx, 'G'] = group_number
    group_number += 1

print(df)

代码解释 #

1. 初始化G列：先创建空的G列，确保所有行初始状态为空。
2. 分组遍历：按A列分组，逐个处理每个分组。
3. 定位目标行：
   • 先找到组内P=1的行索引；
   • 从该索引开始，筛选所有P=0的行，取第二个结果的索引（对应需求中的"P=1之后的第二个0"）。
4. 赋值组编号：给目标行的G列填入当前组的序号，完成后组序号自增，处理下一个分组。

运行结果 #

执行代码后，输出的DataFrame与期望一致：

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