如何在Pandas groupby()分组结果中添加非分组的学生列

原始DataFrame定义 #

import pandas as pd

df = pd.DataFrame({
    'student': ['Bob', 'Sally', 'Rich', 'Melissa', 'Len', 'Sue', 'Jon', 'Sandy',
                'William', 'Sara'],
    'year': [2020, 2020, 2020, 2021, 2021, 2021, 2021, 2022, 2022, 2022],
    'gpa': [2.9, 3.7, 3.2, 3.8, 3.8, 3.1, 3.2, 2.7, 3.6, 3.9]})

输出：

student  year  gpa
0      Bob  2020  2.9
1    Sally  2020  3.7
2     Rich  2020  3.2
3  Melissa  2021  3.8
4      Len  2021  3.8
5      Sue  2021  3.1
6      Jon  2021  3.2
7    Sandy  2022  2.7
8  William  2022  3.6
9     Sara  2022  3.9

用户尝试的代码及问题 #

用户尝试通过groupby获取每年最高GPA，但结果缺少学生列：

result = df.groupby(['year'])[['gpa']].max().reset_index()

输出：

year    gpa
0   2020    3.7
1   2021    3.8
2   2022    3.9

需要得到包含对应学生（含并列）的结果：

year   gpa       student
2020   3.7       Sally
2021   3.8       Len
2021   3.8       Melissa
2022   3.9       Sara

解决方案 #

方法1：先计算年度最高GPA，再筛选原数据 #

先通过groupby得到每年的最高GPA，再用布尔索引匹配原数据中符合条件的行：

# 计算每年的最高GPA
year_max_gpa = df.groupby('year')['gpa'].max().reset_index()
# 合并并筛选
result = df.merge(year_max_gpa, on=['year', 'gpa'])
# 调整列顺序（可选）
result = result[['year', 'gpa', 'student']]

输出：

year  gpa   student
0  2020  3.7     Sally
1  2021  3.8   Melissa
2  2021  3.8       Len
3  2022  3.9      Sara

方法2：使用transform方法（更简洁） #

transform会将分组后的结果广播到原数据的每一行，直接筛选出等于年度最高GPA的行：

# 标记出每行是否等于当年最高GPA
df['is_max_gpa'] = df['gpa'] == df.groupby('year')['gpa'].transform('max')
# 筛选标记为True的行，再去掉辅助列
result = df[df['is_max_gpa']][['year', 'gpa', 'student']]

输出和方法1一致。

方法3：使用rank方法 #

通过排名筛选出当年GPA排名第一的学生（支持并列）：

# 按年份分组，对GPA降序排名，并列的取相同排名
df['gpa_rank'] = df.groupby('year')['gpa'].rank(method='min', ascending=False)
# 筛选排名为1的行
result = df[df['gpa_rank'] == 1][['year', 'gpa', 'student']]

输出同样符合要求。

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