如何基于Group_ID分组的最大值生成argmax新列

现有如下结构的DataFrame：

Group_ID    probability 
34883       0.002676    
34883       0.17826266  
34883       0.01399753  
34883       0.04569782  
34883       0.02799506  
34883       0.02634829  
34883       0.02923014  
34883       0.13544669  
34883       0.07595718  
34883       0.19246604  
34883       0.20028818  
34883       0           
34883       0           
34883       0.07163442  
34897       0.03329843  
34897       0.07643979  
34897       0.09570681  
34897       0.00376963  
34897       0.01780105  
34897       0.0008377   
34897       0.08125654  
34897       0.10764398  
34897       0.25780105  
34897       0.10910995  
34897       0           
34897       0.02743455  
34897       0.18890052  
34897       0           

每个Group_ID对应的probability列数值之和为1。需要创建名为argmax的新列，标记每组中概率最大的行：组内概率最大的行标记为1，否则为0，预期结果如下：

Group_ID    probability argmax
34883       0.002676    0
34883       0.17826266  0
34883       0.01399753  0
34883       0.04569782  0
34883       0.02799506  0
34883       0.02634829  0
34883       0.02923014  0
34883       0.13544669  0
34883       0.07595718  0
34883       0.19246604  1
34883       0.20028818  0
34883       0           0
34883       0           0
34883       0.07163442  0
34897       0.03329843  0
34897       0.07643979  0
34897       0.09570681  0
34897       0.00376963  0
34897       0.01780105  0
34897       0.0008377   0
34897       0.08125654  0
34897       0.10764398  0
34897       0.25780105  1
34897       0.10910995  0
34897       0           0
34897       0.02743455  0
34897       0.18890052  0
34897       0           0

想使用idxmax但不确定具体实现方法，寻求帮助。

这里提供两种高效实现方式，其中包含你提到的idxmax用法：

方法一：分组广播最大值对比 #

利用transform将组内最大值映射到每一行，再通过比较生成标记列：

import pandas as pd

# 假设你的DataFrame名为df
df['argmax'] = (df['probability'] == df.groupby('Group_ID')['probability'].transform('max')).astype(int)

• groupby('Group_ID')['probability'].transform('max')：计算每个分组的概率最大值，并将该值填充到组内所有行，生成与原DataFrame等长的Series。
• 对比每行概率是否等于组内最大值，得到布尔值后用astype(int)转为0/1格式。

方法二：使用idxmax标记目标行 #

先获取每组概率最大值对应的行索引，再批量标记：

# 获取每个分组中概率最大的行索引
max_row_indices = df.groupby('Group_ID')['probability'].idxmax()

# 初始化argmax列为0
df['argmax'] = 0
# 将目标索引行的argmax设为1
df.loc[max_row_indices, 'argmax'] = 1

• groupby('Group_ID')['probability'].idxmax()：返回每个分组中第一个出现概率最大值的行索引。
• 先统一设为0，再通过索引定位赋值1，适合只标记第一个最大值行的场景。

注：如果同一分组存在多个概率相等的最大值，方法一会全部标记为1，方法二则仅标记第一个出现的行，可根据需求选择。

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