如何在Pandas v2中优化分组筛选极值行（含NaN处理）

问题背景 #

给定如下Pandas DataFrame：

import pandas as pd
import numpy as np

x = pd.DataFrame({
    "t": ["A", "A", "A", "B", "B", "B", "B", "A", "C", "C"],
    "m": ["k", "m", "m", "k", "m", "m", "b", "k", "f", "d"],
    "f1": [1.2, np.nan,  0.8, 1, 1, 1.5, 1, np.nan, np.nan, np.nan],
    "f2": [100, 200, 200, 100, 100, 100, 50, 200, 300, 400]})

对应数据：

t  m   f1   f2
0  A  k  1.2  100
1  A  m  NaN  200
2  A  m  0.8  200
3  B  k  1.0  100
4  B  m  1.0  100
5  B  m  1.5  100
6  B  b  1.0   50
7  A  k  NaN  200
8  C  f  NaN  300
9  C  d  NaN  400

需求为：

• 按t列分组，对f1列保留每组最小值对应的行；若组内f1全为NaN，则保留该组所有行
• 基于上述结果，对f2列保留每组最大值对应的行；若组内f2全为NaN，则保留该组所有行

期望输出：

t  m   f1   f2
0  A  m  0.8  200
1  B  k  1.0  100
2  B  m  1.0  100
3  C  d  NaN  400

现有实现代码 #

已通过以下代码实现需求：

def keep_rows(k, col, op):
    # 若组内该列全为NaN则返回整组
    if np.isnan(k[col].values).all():
        return k 
    return k[k[col] == getattr(np, f"nan{op}")(k[col])]

# 先按f1筛选，再按f2筛选
tt = x.groupby("t", as_index=False).apply(lambda x: keep_rows(x, "f1", "min")).reset_index(drop=True)
tt = tt.groupby("t", as_index=False).apply(lambda x: keep_rows(x, "f2", "max")).reset_index(drop=True)

更优实现方法（Pandas v2.x+） #

方法一：Transform + 布尔索引（性能最优） #

利用Pandas的transform方法生成组内统计值的广播列，结合布尔索引实现筛选，避免了apply的循环开销，性能更优，逻辑更直观：

# 第一步：处理f1列筛选
# 生成每组f1的nan最小值广播列
f1_min = x.groupby('t')['f1'].transform(np.nanmin)
# 生成每组f1是否全为NaN的广播列
f1_all_nan = x.groupby('t')['f1'].transform(lambda g: g.isna().all())
# 筛选条件：f1等于组内最小值 或 组内f1全为NaN
mask_f1 = (x['f1'] == f1_min) | f1_all_nan
df_step1 = x[mask_f1]

# 第二步：处理f2列筛选
f2_max = df_step1.groupby('t')['f2'].transform(np.nanmax)
f2_all_nan = df_step1.groupby('t')['f2'].transform(lambda g: g.isna().all())
mask_f2 = (df_step1['f2'] == f2_max) | f2_all_nan
result = df_step1[mask_f2].reset_index(drop=True)

print(result)

方法二：利用Pandas 2.x原生nanmin/nanmax方法 #

Pandas 2.x为GroupBy对象新增了nanmin/nanmax原生方法，无需调用numpy的对应函数，代码更简洁：

# 处理f1筛选
# 获取每组f1的nan最小值
f1_group_min = x.groupby('t')['f1'].nanmin().reset_index(name='f1_min')
# 获取每组f1是否全为NaN
f1_all_nan = x.groupby('t')['f1'].isna().all().reset_index(name='f1_all_nan')
# 合并后筛选
df_f1 = x.merge(f1_group_min, on='t').merge(f1_all_nan, on='t')
df_step1 = df_f1[(df_f1['f1'] == df_f1['f1_min']) | df_f1['f1_all_nan']].drop(columns=['f1_min', 'f1_all_nan'])

# 处理f2筛选
f2_group_max = df_step1.groupby('t')['f2'].nanmax().reset_index(name='f2_max')
f2_all_nan = df_step1.groupby('t')['f2'].isna().all().reset_index(name='f2_all_nan')
df_f2 = df_step1.merge(f2_group_max, on='t').merge(f2_all_nan, on='t')
result = df_f2[(df_f2['f2'] == df_f2['f2_max']) | df_f2['f2_all_nan']].drop(columns=['f2_max', 'f2_all_nan']).reset_index(drop=True)

print(result)

方案对比 #

• 现有实现：依赖apply循环处理每组，在数据量较大时性能较差，代码嵌套程度高
• 更优方案：
  • Transform方式利用Pandas向量化操作，避免循环，性能提升明显
  • Pandas 2.x原生方法让代码更简洁，无需依赖numpy的nan函数，兼容性更好

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