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

Input:

    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，label 2。现在，我想使用列名最高概率的label作为id的预测值。如果有两个标签的可能性最高，请选择其中一个。

Expected:

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

谢谢！

Answer 1

对每行（即axis = 1）使用idxmax：

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

Answer 2

解决方案：

a['prediction']=a[['label 0','label 1','label 2']].apply(max,axis=1)

演示：

a['prediction']=a[['label 0','label 1','label 2']].apply(max,axis=1)
a=a.drop(['label 0','label 1','label 2'],1)
print(a)

    id  prediction
0   1         0.5
1   2         0.4
2   3         0.5