Numpyで1-of-K表現を簡単に作る方法

愚直にやるとめんどくさそうでしたが、結構簡単に出来たのでメモ。 一応、1-of-K表現とは、ある要素だけが1で、それ以外が0なベクトルです。

[0, 0, 1, 0, 0]

みたいな奴。

ここでは、あるidxの要素が1、それ以外が0であることがわかっていると仮定する。

import numpy as np
K = 10
idx = 2
a = np.zeros(K)
a.put(idx, 1) # ==> array([ 0.,  0.,  1.,  0.,  0.,  0.,  0.,  0.,  0.,  0.])

簡単！