根据来自另一个数组 numpy 的索引分配值

Question

我有一个这样的索引数组：

idx = np.array([3,4,1], [0,0,0], [1,4,1], [2,0,2]]

还有一个零数组A，形状为4x5

我想将A 的idx 中的所有索引设为1

对于上面的例子，最终的数组应该是：

[[0,1,0,1,1],  # values at index 3,4,1 are 1
 [1,0,0,0,0],  # value at index 0 is 1
 [0,1,0,0,1],  # values at index 1,4 are 1
 [1,0,1,0,0]]  # values at index 0,2 are 1

如何在 numpy 中做到这一点？

Answer 1

使用花哨的索引：

A = np.zeros((4, 5), dtype=int)

A[np.arange(len(idx))[:,None], idx] = 1

更新A:

array([[0, 1, 0, 1, 1],
       [1, 0, 0, 0, 0],
       [0, 1, 0, 0, 1],
       [1, 0, 1, 0, 0]])