Pandas DataFrame分组去重并保留cnt列最大值行

问题场景 #

你有如下结构的DataFrame，需要按name、size、yr分组，每组内保留cnt数值最大的对应行：

import pandas as pd
import numpy as np

df = pd.DataFrame({
                   'name': ['p1', 'p1', 'p2', 'p2'],
                   'size': [10, 10, 20, 20],
                   'yr': [1990, 1990, 2000, 2000],
                   'res': [1, 2, 1, 2],
                   'cnt': [100, 80, 220, 400]
                  })

期望输出：

name  size    yr  res  cnt
0   p1    10  1990    1  100
3   p2    20  2000    2  400

解决方案 #

方法1：使用groupby + idxmax

先按分组列（name、size、yr）分组，定位每组cnt最大值对应的行索引，再提取目标行：

# 定义分组依据的列
group_keys = ['name', 'size', 'yr']
# 获取每组cnt最大值的行索引
max_cnt_idx = df.groupby(group_keys)['cnt'].idxmax()
# 提取对应行并重置索引
result = df.loc[max_cnt_idx].reset_index(drop=True)
print(result)

方法2：使用sort_values + drop_duplicates

先按分组列升序、cnt降序排序，再按分组列去重并保留第一行（即cnt最大的行）：

# 排序：分组列升序，cnt降序
sorted_df = df.sort_values(by=[*group_keys, 'cnt'], ascending=[True, True, True, False])
# 按分组列去重，保留第一个（cnt最大的）
result = sorted_df.drop_duplicates(subset=group_keys, keep='first').reset_index(drop=True)
print(result)

两种方法均能得到目标结果：方法1在大数据量下效率更高，方法2逻辑更直观易懂。

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