Как выбрать имя столбца с наибольшим значением в качестве нового элемента столбца?

Question

У меня есть DataFrame с именем a, который может быть показан следующим образом.

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

Решение:

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

Answer 2

Используйте idxmax для каждой строки (т. Е. Ось = 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