如何在Python中用idxmax()选取列值相同时的最大索引行？

问题场景 #

现有如下DataFrame数据：

pcn_no   treatment   score
1   101    3000        0.005804
2   101    5000        0.005804
3   101    8000        0.005804
5   101    12000       0.005804
8   101    25000       0.005804
12  101    50000       0.005804
4   101    10000       0.014521
6   101    15000       0.014521
7   101    20000       0.014521
9   101    30000       0.014521
10  101    35000       0.014521
11  101    40000       0.014521

需求是：按pcn_no分组后，选取每组中score最大值对应的行里索引最大的那一行（即示例中treatment为40000的行），但使用以下代码时，返回的却是该最大值组中索引最小的行（treatment为10000）：

x = 'score'
df_all[x + '_max_line'] = df_all.loc[(df_all.groupby('pcn_no')[x].idxmax(),'treatment')]

问题原因 #

idxmax()函数默认返回分组中第一个出现目标列最大值的索引，因此在score值相同的组里，会优先取索引最小的行，不符合需求。

解决方案 #

方法1：排序后取分组首行 #

先对整个DataFrame按score降序、索引降序排序，再按pcn_no分组取每组第一行，即可得到每组中score最大且索引最大的行：

# 排序：先按score降序，再按索引降序
sorted_df = df_all.sort_values(['score', df_all.index], ascending=[False, False])
# 分组取首行
max_rows = sorted_df.groupby('pcn_no').first()

# 将结果映射回原DataFrame
df_all['score_max_line'] = df_all['pcn_no'].map(max_rows['treatment'])

方法2：自定义分组逻辑获取目标索引 #

针对每个分组，先筛选出score等于最大值的所有行，再取这些行的最大索引：

def get_target_idx(group):
    max_score = group['score'].max()
    # 筛选score等于最大值的行，返回最大索引
    return group[group['score'] == max_score].index.max()

# 获取每个pcn_no对应的目标索引
target_indices = df_all.groupby('pcn_no').apply(get_target_idx)

# 提取对应treatment值并合并到原DataFrame
df_all['score_max_line'] = df_all.loc[target_indices, 'treatment'].reindex(df_all['pcn_no']).values

方法3：利用负索引反向取最大值 #

由于idxmax()会返回第一个最小值的索引，我们可以对索引取负数，将原最大索引转换为最小的负数，再用idxmax()获取：

# 分组后筛选score最大值的行，取负索引的idxmax（等价于原索引的最大值）
target_indices = df_all.groupby('pcn_no').apply(
    lambda g: g[g['score'] == g['score'].max()].index * -1).idxmax()

# 映射treatment值到原DataFrame
df_all['score_max_line'] = df_all.loc[target_indices, 'treatment'].reindex(df_all['pcn_no']).values

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