如何按条件筛选Pandas DataFrame？满足特定距离筛选规则

原始数据 #

import pandas as pd

df = pd.DataFrame({
    'city': ['New York','New York','New York','Los Angeles','Los Angeles','Houston','Houston','Houston'],
    'airport': ['LGA', 'EWR', 'JFK', 'LAX', 'BUR', 'IAH', 'HOU', 'EFD'],
    'distance': [38, 50, 32, 8, 50, 90, 78, 120]
})

筛选需求 #

• 若某城市存在距离≤40的机场，保留所有满足该条件的行
• 若某城市没有距离≤40的机场，仅保留该城市中距离最短的一行

实现代码 #

# 标记每个城市是否有距离≤40的机场
df['has_near'] = df.groupby('city')['distance'].transform(lambda x: (x <= 40).any())

# 筛选目标数据
result = df[
    (df['distance'] <= 40) | 
    (~df['has_near'] & (df['distance'] == df.groupby('city')['distance'].transform('min')))
].drop(columns='has_near')

print(result)

输出结果 #

city airport  distance
0    New York     LGA        38
2    New York     JFK        32
3  Los Angeles     LAX         8
6      Houston     HOU        78

逻辑说明 #

• 用groupby.transform给每个城市的所有行标记是否存在近机场（距离≤40），生成临时辅助列has_near
• 筛选条件分两类：
  1. 直接保留距离≤40的行
  2. 针对无近机场的城市，筛选出该城市距离等于最小值的行
• 最后删除临时辅助列，得到最终结果

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