如何为Pandas DataFrame新增列存储概率最高的类别？

嘿，你已经走对方向啦！你用agg(max)拿到了每行的最大概率值，现在只需要把这个方法换成idxmax()，就能直接得到对应概率最高的类别名称了。

具体实现代码 #

import pandas as pd

# 初始化数据
dict_test = {'series': [1, 2, 3, 4, 5, 6, 7], 
             'cat_1': [.02, .02, .81, .72, .01, .3, .45], 
             'cat_2': [.02, .02, .14, .2, .99, .45, .4], 
             'cat_3': [.96, .96, .05, .08, .00, .25, .15]}
df = pd.DataFrame(dict_test)

# 新增result列，获取每行概率最高的类别
df['result'] = df.drop('series', axis=1).idxmax(axis=1)

# 查看最终结果
print(df)

为什么这个方法有效？ #

idxmax(axis=1)会遍历DataFrame的每一行（axis=1指定按行操作），返回该行中数值最大的列标签——正好就是你需要的cat_1/cat_2/cat_3这些类别名称，完美匹配你的期望结果。

补充说明：并列最大值的情况 #

如果某一行有多个类别概率相同且都是最大值，idxmax()会返回第一个出现的类别。比如假设某行cat_1和cat_2概率都是0.5，结果会是cat_1，因为它在列顺序中更靠前。

运行上面的代码后，你的DataFrame就会和你想要的结果完全一致啦！

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