如何在Python的Pandas DataFrame中提取各分组列最大值的所有行？

问题描述 #

有一个包含col1、col2、col3、col4列的Pandas DataFrame，需要提取每个col1分组中col2为最大值的所有行。

示例数据 #

col1  col2  col3   col4
A101  3     LLT    10028980
A101  7     LLT    10028980
A101  7     PT     10028980
A102  5     LLT    10028981
A102  3     PT     10028981 
A103  2     PT     10028982
A103  4     LLT    10028982

期望输出 #

col1  col2  col3   col4
A101  7     LLT    10028980
A101  7     PT     10028980
A102  5     LLT    10028981
A103  4     LLT    10028982

原代码问题 #

原代码使用groupby('col1').max()会直接取每个分组各列的最大值，过滤掉了同一分组中col2为最大值的多行记录：

m = df.notnull().all(axis=1)
df = df.loc[m].groupby('col1').max().reset_index()

得到的结果丢失了A101分组中col2=7的其中一行。

解决方案 #

可以通过以下几种方法实现需求：

方法1：使用transform匹配最大值 #

利用groupby.transform计算每个分组的col2最大值，然后筛选出col2等于该最大值的行：

# 先过滤空行（保留原逻辑）
m = df.notnull().all(axis=1)
df_clean = df.loc[m]

# 计算每个分组的col2最大值
max_col2_per_group = df_clean.groupby('col1')['col2'].transform('max')

# 筛选符合条件的行
result = df_clean[df_clean['col2'] == max_col2_per_group]

方法2：使用groupby.filter #

通过filter方法直接筛选出每个分组中col2等于该分组最大值的行：

m = df.notnull().all(axis=1)
df_clean = df.loc[m]

result = df_clean.groupby('col1').filter(lambda x: x['col2'] == x['col2'].max())

方法3：使用rank #

给每个分组内的col2按降序排名，筛选排名为1的行（支持并列最大值）：

m = df.notnull().all(axis=1)
df_clean = df.loc[m]

df_clean['rank'] = df_clean.groupby('col1')['col2'].rank(ascending=False, method='min')
result = df_clean[df_clean['rank'] == 1].drop('rank', axis=1)

以上三种方法都能保留同一分组中col2为最大值的所有行，得到符合期望的输出。

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