python - 如何选择数值最大的列名作为新的列元素？

Question

我有一个名为a的DataFrame，如下所示。

import pandas as pd 
a=pd.DataFrame({'id':[1,2,3],'label 0':[0.2,0.4,0.5],
'label 1':[0.3,0.2,0.1],'label 2':[0.5,0.4,0.4]})

输入:

    id  label 0  label 1  label 2
0   1     0.2     0.3     0.5
1   2     0.4     0.2     0.4
2   3     0.5     0.1     0.4

每一行包含每个id属于三个标签的概率，其中三个标签是label 0、label 1、标签2。我现在想使用概率最高的label(列名)作为id的预测值。如果有两个概率最高的标签，则取其中之一。

预期:

    id   predict  
0   1    label 2     
1   2    label 0     
2   3    label 0 

提前致谢!

Answer 1

使用idxmax对于每一行(即 axis=1):

a.drop('id', 1).idxmax(1)

#0    label 2
#1    label 0
#2    label 0

concat 与 id 列:

pd.concat([a.id, a.drop('id', 1).idxmax(1).rename('predict')], 1)

#   id  predict
#0   1  label 2
#1   2  label 0
#2   3  label 0