argsortの逆関数

概要

numpy.argsort関数の「逆関数」みたいな操作をしたいと思うことがあり、ネットを調べてみても、すぐには出てこなさそうだったので、自分で考えてみました。
いろいろ試した結論としては、numpy.argsort( numpy.argsort( ARRAY_INPUT ) ) とするとよさそうです。

・・・という記事です。

本当？

と思われる方は、以下本文を参照ください・・・

なぜそんなことを考えたのか？

機械学習の仕事をしていて、学習モデルが出力する予測値を変換してからスコア算出する、ということをしたいことがありました。
例えば、学習モデルが以下のような予測値の配列を出力したとします。

INDEX	0	1	2	3	4	5	6	7	8	9
VALUE	0.15	0.04	0.20	0.08	0.05	0.30	0.12	0.01	0.03	0.02

これは0-9の各クラスのいずれかに属するかを示す確信度（1に近いほど高確率）を表しているものとします。このままでもいいのかもしれないですが、トップ3だけに、順に10点、5点、1点のスコアを与え、あとは0点にしてスコア計算したいという要求がありました。
（そういうスコアの与え方でGridSearchしたかったのです。）

つまり、上記の例でいうと、以下のような配列に変換してから、スコア計算したかったわけです。

INDEX	0	1	2	3	4	5	6	7	8	9
VALUE	1	0	5	0	0	10	0	0	0	0

この配列をどうやったら作れるんだろう？
numpy.argsort関数の逆関数、例えばinv_argsortという名前の関数があって、以下のようなことができたらうれしいわけです。

In [1]: prediction
Out[1]: [0.15, 0.04, 0.2, 0.08, 0.05, 0.3, 0.12, 0.01, 0.03, 0.02]

In [2]: score = [0, 0, 0, 0, 0, 0, 0, 1, 5, 10]

In [3]: prediction_scored = numpy.take(score, inv_argsort(prediction))

In [4]: prediction_scored
Out[4]: array([ 1,  0,  5,  0,  0, 10,  0,  0,  0,  0])

inv_argsortの実装

概要にも書きましたが、以下の通り、実装自体は簡単です。

def inv_argsort(array_input):
    return numpy.argsort(numpy.argsort(array_input))

でも本当にそれでいいの？という気もしますので、以下で簡単な証明をしておこうと思います。

証明

1回目のargsort

${ n }$個の要素を持つ配列 ${ v }$ が学習モデルの直接の出力とします。

|INDEX| 0 | 1 | 2 | ${\dots}$ | ${n-1}$ |
|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
|${ v }$の値|${ v_0 }$|${ v_1 }$|${ v_2 }$| ${ \dots }$| ${ v_{n-1} }$|

これを numpy.argsort にかけた結果を ${ u }$ とします。

|INDEX| 0 | 1 | 2 | ${\dots}$ | ${n-1}$ |
|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
|${ u }$の値|${ u_0 }$|${ u_1 }$|${ u_2 }$| ${ \dots }$| ${ u_{n-1} }$|
|${ v }$の値|${ v_{u_0} }$|${ v_{u_1} }$|${ v_{u_2} }$| ${ \dots }$| ${ v_{u_{n-1}} }$|

ここで、argsortの定義により、

$$ v_{u_0} \leq v_{u_1} \leq v_{u_2} \leq \dots \leq v_{u_{n-1}} $$

であるとわかります。
つまり、INDEXの $0, 1, 2,\dots, n-1$ と大小関係が一致するわけです。
ここでこの表の最下行に、付与したいスコアもくっつけてみます。

|INDEX| 0 | 1 | 2 | ${\dots}$ | ${n-1}$ |
|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
|${ u }$の値|${ u_0 }$|${ u_1 }$|${ u_2 }$| ${ \dots }$| ${ u_{n-1} }$|
|${ v }$の値|${ v_{u_0} }$|${ v_{u_1} }$|${ v_{u_2} }$| ${ \dots }$| ${ v_{u_{n-1}} }$|
|SCORE|${ s_0 }$|${ s_1 }$|${ s_2 }$| ${ \dots }$|${ s_{n-1} }$|

概要の例でいえば、${ s_9 = 10, s_8 = 5, s_7 = 1 }$ で残りは 0 ということになります。
$v$ の昇順に並んでいるので、「トップ3に順に 10点、5点、1点を割り当てる」という対応関係が成立しています。

2回目のargsort

$u$ を numpy.argsort にかけた結果を ${ w }$ とします。

|INDEX| 0 | 1 | 2 | ${\dots}$ | ${n-1}$ |
|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
|${ w }$の値|${ w_0 }$|${ w_1 }$|${ w_2 }$| ${ \dots }$| ${ w_{n-1} }$|
|${ u }$の値|${ u_{w_0} }$|${ u_{w_1} }$|${ u_{w_2} }$| ${ \dots }$| ${ u_{w_{n-1}} }$|
|${ v }$の値|${ v_{u_{w_0}} }$|${ v_{u_{w_1}} }$|${ v_{u_{w_2}} }$| ${ \dots }$| ${ v_{u_{w_{n-1}}} }$|
|SCORE|${ s_{w_0} }$|${ s_{w_1} }$|${ s_{w_2} }$| ${ \dots }$|${ s_{w_{n-1}}} $|

ここで、あることに気が付きます。
$$ u_{w_i} \in [0, 1, 2, \dots, n-1], \hspace{12pt} u_{w_0} \lt u_{w_1} \lt u_{w_2} \lt \dots \lt u_{w_{n-1}}$$
という事実から、
$$ u_{w_0}=0, u_{w_1}=1, u_{w_2}=2, \dots , u_{w_{n-1}}=n-1 $$
であるとわかります。
これをもとに表を書き換えておきましょう。

|INDEX| 0 | 1 | 2 | ${\dots}$ | ${n-1}$ |
|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
|${ w }$の値|${ w_0 }$|${ w_1 }$|${ w_2 }$| ${ \dots }$| ${ w_{n-1} }$|
|${ u }$の値|${ 0 }$|${ 1 }$|${ 2 }$| ${ \dots }$|${ n-1 }$|
|${ v }$の値|${ v_0 }$|${ v_1 }$|${ v_2 }$| ${ \dots }$| ${ v_{n-1} }$|
|SCORE|${ s_{w_0} }$|${ s_{w_1} }$|${ s_{w_2} }$| ${ \dots }$|${ s_{w_{n-1}}} $|

・・・おや、$v$ の並び順がもともとの並び順に戻っていますね。
これはつまり、学習モデルが出力したままの順序に戻ったわけです。
それに対応するスコア配列の並びは、 $w$ に記録されています。

つまり、求めるスコア配列（prediction_scored）は、以下で求めることができるということです。

v = prediction
u = numpy.argsort( v )
w = numpy.argsort( u )
score = [0, 0, 0, 0, 0, 0, 0, 1, 5, 10]
prediction_scored = numpy.take(score, w)

[証明終了]