如何在Pandas分组聚合中获取最小值及最小值+1的行？

需求说明 #

按type字段分组，筛选每个组内distance等于组内最小值或最小值+1的行。

示例数据 #

import pandas as pd

df = pd.DataFrame({
    'distance': [8, 1, 1, 2, 3, 3],
    'vertex': [104, 114, 103, 102, 18, 108],
    'type': ['A', 'A', 'B', 'A', 'A', 'B']
})

现有代码局限 #

你当前的代码仅筛选出了组内最小值的行，无法覆盖最小值+1的情况：

mask = df['distance'].isin(df.groupby(['type'])['distance'].agg('min').values)
df[mask]

解决方案 #

使用transform方法将分组计算的最小值广播到每一行，然后直接判断distance是否符合条件：

# 计算每个组的最小值，并映射到对应组的每一行
group_min = df.groupby('type')['distance'].transform('min')

# 构建筛选条件：distance等于最小值 或 最小值+1
mask = (df['distance'] == group_min) | (df['distance'] == group_min + 1)

# 获取结果
result = df[mask]
print(result)

输出结果 #

distance  vertex type
1         1     114    A
2         1     103    B
3         2     102    A

方法说明 #

• transform('min')会为每个组计算最小值，并将该值填充到组内的每一行，保证每行都能拿到对应组的最小值。
• 直接通过向量化的条件判断（==和|）筛选行，比循环或apply更高效，适合处理大数据集。

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