选择最高价值列

Question

我需要你的帮助。数据框存储以下三种类别的概率：

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]}

import pandas as pd
df = pd.DataFrame(dict_test)

我需要创建一个新列来存储哪个类别具有最高概率。到目前为止，我已经能够使用agg函数选择最高的概率：

df['choice'] = df.drop('series', axis = 1).agg(max, axis = 1)

我需要的结果以该数据框为例：

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],
             'result': ['cat_3', 'cat_3', 'cat_1', 'cat_1', 'cat_2', 'cat_2', 'cat_1']}

df = pd.DataFrame(dict_test)

有什么建议吗？

Answer 1

您正在寻找idxmax

df['result'] = df.filter(regex='^cat').idxmax(axis=1)