python - 索引位置为 1 的数组

Question

我目前有一个数组中最小值的索引数组。

看起来像这样:

[[0],
 [1],
 [2],
 [1],
 [0]]

(最大索引为3)

我想要的是一个如下所示的数组:

[[1, 0, 0]
 [0, 1, 0]
 [0, 0, 1]
 [0, 1, 0]
 [1, 0, 0]]

其中 1 在最小值的列中。

在 numpy 中有没有简单的方法来做到这一点？

Answer 1

使用 NumPy 的 == 广播:

>>> minima = np.array([[0], [1], [2], [1], [0]])
>>> minima == arange(minima.max() + 1)
array([[ True, False, False],
       [False,  True, False],
       [False, False,  True],
       [False,  True, False],
       [ True, False, False]], dtype=bool)
>>> (minima == arange(minima.max() + 1)).astype(int)
array([[1, 0, 0],
       [0, 1, 0],
       [0, 0, 1],
       [0, 1, 0],
       [1, 0, 0]])