Python
numpy
pandas

Pandasで数値の大きな順に列番号を取り出す

多クラス分類や多ラベル分類などで目的変数ごとに複数の確率値が出てきた場合、確率値の高い順に並べたいということはよくあるので、Pandasでのやり方を簡単に記載しておきます。

まず例として列ごとに各クラスの確率値が入ったDataFrameを作成します。ちなみにここで各クラス名を0,1,2にしているのは、このあと作成するインデックスの数字に合わせているためです。後ほど変換できるので、任意の名前で問題ありません。

In [1]: import numpy as np

In [2]: import pandas as pd

In [3]: pred = pd.DataFrame({"class0":[0.1,0.4,0.3], "class1":[0.5,0.2,0.4], "class2":[0.3,0.5,0.8]})

In [4]: pred
Out[4]: 
   class0  class1  class2
0     0.1     0.5     0.3
1     0.4     0.2     0.5
2     0.3     0.4     0.8

ここでnp.argsortを使ってインデックスを貼ります。並び方が昇順になってしまうので[::-1]で降順になおしています。DataFrame自体のカラム名が変になりますが気にしない。

In [5]: pred_rank = pred.apply(lambda x: np.argsort(x)[::-1], axis=1)

In [6]: pred_rank
Out[6]: 
   class2  class1  class0
0       1       2       0
1       2       0       1
2       2       1       0

例えば0行目はpredを見れば値を見ればclass1→class2→class0の順なので、きちんと左のカラムから順に大きなクラスの値が並んでいるのが分かると思います。あとは煮るなり焼くなり好きにします。

確率値の高い上位n位を取り出す

ここでは最も確率値の高いクラスを表示しています。複数の場合はx[:1]部分を任意のnに置き換えればOKです。

In [7]: pred_rank.apply(lambda x: x[:1], axis=1)
Out[7]: 
   class2
0       1
1       2
2       2

クラス名を付けて整形する

インデックスだけではよく分からないので、クラス名を付与して表示します。

In [8]: pred_rank.apply(lambda x: " -> ".join(pred.columns[x]), axis=1)
Out[8]: 
0    class1 -> class2 -> class0
1    class2 -> class0 -> class1
2    class2 -> class1 -> class0
dtype: object

参考