DataFrame特定条件过滤：保留最后一个value>500及之后的行

原始DataFrame #

Id      action          date          value

A       enter          20/12/2021       0
A       enter          20/12/2021      150
A       L-1            20/12/2021      520
A       L-2            20/12/2021      50
A       L-2            20/12/2021      550
A       L-3            20/12/2021      20
A       L-4            20/12/2021      5
A       L-5            20/12/2021      1
B       enter          25/12/2021      2
B       L-1            25/12/2021      510
B       L-2            25/12/2021      6
B       L-3            25/12/2021      3
C       enter          26/12/2021      4
C       L-1            26/12/2021      10
C       L-2            26/12/2021      20

预期结果 #

Id      action          date          value

A       L-2            20/12/2021      550
A       L-3            20/12/2021      20
A       L-4            20/12/2021      5
A       L-5            20/12/2021      1
B       L-1            25/12/2021      510
B       L-2            25/12/2021      6
B       L-3            25/12/2021      3
C       enter          26/12/2021      4
C       L-1            26/12/2021      10
C       L-2            26/12/2021      20

需求说明 #

• 按Id分组，过滤掉每组中最后一个value>500的行之前的所有行，保留该行及其之后的行
• 若分组内无value>500的行，则保留该组全部行

现有代码问题 #

现有代码逻辑存在缺陷，会漏掉符合value>500条件的目标行，代码如下：

g = df.loc[::-1, 'value'].le(500).groupby(df['Id'])

m1 = g.cummin()           # drop rows with >500 and before
m2 = ~g.transform('all')  # drop groups without value >500

out = df.loc[m1&m2]

修正后的代码 #

方法一：基于索引定位 #

# 找到每个分组中最后一个value>500的行的索引
last_gt500_idx = df[df['value'] > 500].groupby('Id')['value'].idxmax()

# 生成掩码：有符合条件的组保留从目标行开始的所有行；无则全保留
mask = df.index.to_series().apply(lambda x: x >= last_gt500_idx.get(df.loc[x, 'Id'], x))

# 输出结果
out = df[mask]

方法二：基于分组累积标记 #

# 反转DataFrame后，从后往前标记需要保留的行（找到第一个value>500后，后续所有行都标记为保留）
df['keep'] = df[::-1].groupby('Id')['value'].apply(lambda x: x.gt(500).cumsum().ge(1))

# 给无value>500的分组全量标记为保留
has_gt500 = df.groupby('Id')['value'].max() > 500
df.loc[df['Id'].isin(has_gt500[~has_gt500].index), 'keep'] = True

# 生成最终结果并清理临时列
out = df[df['keep']].drop('keep', axis=1)

两种方法均可实现需求，保留最后一个value>500的行及后续内容，同时完整保留无符合条件分组的所有数据。

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