按列值分组并返回指定列最大值及对应其他列数据

针对按列分组后获取指定列最大值所在行全部数据的需求，这里提供几个通用方案，适配任意数量的列：

方法一：用idxmax定位最大行索引 #

先通过groupby找到每组指定列（比如X1）最大值对应的行索引，再用loc提取整行数据：

import pandas as pd
dat = pd.DataFrame({'A' : ['Y', 'Y', 'Z'] , 'X1' : [10, 20, 30], 'X2' : [11, 10, 20], 'X3' : [20, 15, 60]})

# 获取每组X1最大值对应的行索引
max_row_indices = dat.groupby('A')['X1'].idxmax()
# 提取对应行的所有列数据
result = dat.loc[max_row_indices].reset_index(drop=True)
print(result)

输出结果：

A  X1  X2  X3
0  Y  20  10  15
1  Z  30  20  60

方法二：排序后取每组首行 #

先按指定列降序排序，再分组取每组第一行，自然拿到该列最大值所在的行：

result = dat.sort_values('X1', ascending=False).groupby('A').first().reset_index()
print(result)

这个方法逻辑简单，同样支持任意列数的DataFrame。

方法三：保留所有最大值行（若组内有多个最大值） #

如果组内存在多个行的指定列值同为最大值，前两种方法只会取其中一行。若需要保留所有符合条件的行，可以用transform标记组内最大值，再筛选：

# 给每行标记所在组的X1最大值
dat['group_x1_max'] = dat.groupby('A')['X1'].transform('max')
# 筛选出X1等于组内最大值的行
result = dat[dat['X1'] == dat['group_x1_max']].drop('group_x1_max', axis=1).reset_index(drop=True)
print(result)

比如若Y组有两个X1=20的行，这个方法会把两行都保留下来。

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