如何仅依据每组首行数值按条件过滤Pandas分组？

问题描述 #

给定如下DataFrame：

import pandas as pd
df = pd.DataFrame(
    {
        'group': list('xxxxyyy'),
        'open': [100, 150, 200, 160, 300, 150, 170],
        'close': [105, 150, 200, 160, 350, 150, 170],
        'stop': [104, 104, 104, 104, 400, 400, 400]
    }
)

需要过滤出**每组首行的stop值在首行open和close值之间（不考虑顺序）**的所有组，预期输出仅保留group为x的所有行：

group  open  close  stop
0     x   100    105   104
1     x   150    150   104
2     x   200    200   104
3     x   160    160   104

优化解决方案 #

原方法使用groupby.apply逐组判断，虽然可行但在大数据集上效率较低。推荐先筛选出符合条件的分组，再过滤原DataFrame，利用pandas矢量化操作提升性能：

步骤1：提取每组首行数据

group_first = df.groupby('group').first()

步骤2：筛选符合条件的分组

判断首行stop是否在open和close之间（自动处理大小顺序）：

valid_groups = group_first[
    group_first['stop'].between(
        group_first[['open', 'close']].min(axis=1),
        group_first[['open', 'close']].max(axis=1)
    )
].index

步骤3：过滤原DataFrame

out = df[df['group'].isin(valid_groups)]

替代方案（基于transform） #

如果需要保留分组内的判断标记，可以用transform实现：

def check_group_validity(group):
    s = group['stop'].iloc[0]
    o = group['open'].iloc[0]
    c = group['close'].iloc[0]
    return min(o, c) <= s <= max(o, c)

# 给每行标记所属组是否符合条件
df['is_valid'] = df.groupby('group')['group'].transform(check_group_validity)
# 过滤出有效组的行
out = df[df['is_valid']].drop('is_valid', axis=1)

方案对比 #

• 第一种方案（先筛选分组再过滤）：仅处理每组首行，矢量化操作效率最高，代码简洁，适合大数据集。
• 原apply方案：逐组返回数据，会产生额外的数据拼接开销，性能不如前两种方案。

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