如何清洗DataFrame：按A分组，组内存在B>100时仅保留对应行

需求规则 #

按A列分组：

• 若组内存在B列值大于100的行，仅保留该组中B列值>100的行
• 若组内没有B列值>100的行，保留该组所有行
  注意：A列可为非数字类型，原sort_values结合>100的方法因NaN值无法正常工作

输入示例 #

A   B     C...(other columns)
1   100
1   98
2   150
2   100 #(removed)
2   90 #(removed)
3   200
3   75 #(removed)
4   80
4   80
4   70
4   70

输出示例 #

A   B    C...(other columns)
1   100
1   98
2   150
3   200
4   80
4   80
4   70
4   70

可复现代码 #

import pandas as pd
data1 = {'A': [1, 1, 2, 2, 2, 3, 3, 4, 4, 4, 4], 
         'B': [100, 98, 150, 100, 90, 200, 75, 80, 80, 70, 70]}
df = pd.DataFrame(data1)

解决方案 #

用groupby结合transform标记符合条件的组，再进行过滤，该方法能处理NaN值且不限制A列类型：

# 标记每个组是否存在B>100的行
df['has_over_100'] = df.groupby('A')['B'].transform(lambda x: (x > 100).any())

# 执行过滤：要么组内无B>100的行，要么当前行B>100
filtered_df = df[(~df['has_over_100']) | (df['B'] > 100)].drop('has_over_100', axis=1)

print(filtered_df)

该方案的优势：

• 兼容A列的任意数据类型（字符串、数字均可）
• 自动忽略NaN值，不会因空值导致判断逻辑出错
• 无需排序，避免sort_values带来的额外问题

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