如何高效筛选出满足指定mask条件行数≥2的Pandas分组？

原始DataFrame #

import pandas as pd
df = pd.DataFrame({
    'a': [10, 20, 30, 50, 50, 50, 4, 100],
    'b': [30, 3, 200, 25, 24, 31, 29, 2],
    'd': list('aaabbbcc')
})

期望输出 #

a    b  d
0  10   30  a
1  20    3  a
2  30  200  a

需求说明 #

按列d进行分组，筛选出满足**b > a的行数≥2**的分组。

现有实现方案 #

out = df.groupby('d').filter(lambda x: len(x.loc[x.b.gt(x.a)]) >= 2)

更优高效实现方法 #

下面是性能更优的向量化实现方式，尤其适合大数据量场景：

# 1. 计算全局掩码：标记每行是否满足b > a
mask = df['b'] > df['a']
# 2. 统计每个分组中满足条件的行数，筛选出行数≥2的分组
valid_groups = df[mask].groupby('d').size() >= 2
# 3. 获取符合条件的分组标签
valid_d_values = valid_groups[valid_groups].index
# 4. 从原始DataFrame中筛选出属于有效分组的行
out = df[df['d'].isin(valid_d_values)]

优势说明 #

• 避免了filter方法中逐组执行lambda的循环操作，全部采用Pandas向量化API，性能提升显著
• 仅计算一次全局掩码，减少重复计算开销
• 分组统计和筛选逻辑清晰，可读性更强

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