Pandas DataFrame按ID提取连续行至首个flag为1行（含）的实现问题

需求说明 #

现有pandas DataFrame包含id、month、value_1、flag四列，具体处理规则如下：

• 按id分组，每个分组内提取从首行开始到第一个flag=1的所有行（包含flag=1的行）
• 若某id对应所有行flag均为0，提取该id全部行
• 保留原有月份的排列顺序

测试数据集 #

import pandas as pd

data = {'id':['123', '123', '123', '123', '123', '456', '456', '456', '456', '789', '789', '789', '789', '789', '789'],
        'month':[1,2,3,4,5,1,2,3,4,1,2,3,4,5,6],
        'value_1': [232,432,556,223,643,556,121,853,343,324,654,765,128,543,776],
        'flag':[0,0,0,1,1,0,0,0,0,0,0,0,0,1,1]}
 
d = pd.DataFrame(data)

实现代码 #

之前用groupby聚合的方式只能获取单列首个值，无法保留符合要求的全量行，可直接对每个分组做切片处理：

res = d.groupby('id', group_keys=False).apply(
    lambda g: g if g['flag'].max() == 0 else g.loc[:g['flag'].idxmax()]
).reset_index(drop=True)

输出结果 #

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