如何按条件筛选DataFrame子集：保留同category最高note行及NA行

示例数据 #

先定义示例DataFrame：

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'category': ['A', 'A', 'B', 'B', np.nan, np.nan],
    'note': [85, 92, 78, 88, 90, 80],
    'other_col': ['x1', 'x2', 'y1', 'y2', 'z1', 'z2']
})

预期结果 #

expected_df = pd.DataFrame({
    'category': ['A', 'B', np.nan, np.nan],
    'note': [92, 88, 90, 80],
    'other_col': ['x2', 'y2', 'z1', 'z2']
})

高效实现方案 #

方案1：保留同一category下所有note最大值的行

这种方法利用groupby.transform直接生成与原DataFrame长度一致的最大值序列，再通过布尔索引筛选，避免了合并操作，效率更高：

# 生成每组的note最大值序列（非NA组有效，NA组会返回NaN）
max_note = df.groupby('category')['note'].transform('max')
# 筛选条件：category为NA 或 note等于组内最大值
filtered_df = df[(df['category'].isna()) | (df['note'] == max_note)]
# 可选：重置索引
filtered_df = filtered_df.reset_index(drop=True)

方案2：同一category仅保留一行note最大值的行

如果需要每组仅保留一条最大值记录，可使用idxmax获取最大值对应的索引，再拼接NA行：

# 获取非NA组中note最大值的行索引
max_row_indices = df[df['category'].notna()].groupby('category')['note'].idxmax()
# 拼接NA行和选中的最大值行，可按原索引排序保持顺序
filtered_df = pd.concat([df[df['category'].isna()], df.loc[max_row_indices]]).sort_index()

方案优势 #

• 两种方案都避免了冗余的合并操作，处理大数据集时性能更优；
• 逻辑清晰，直接通过布尔索引或索引定位实现筛选，代码简洁易维护。

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