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

Question

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

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

一个形状为4x5的零数组A

我想让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中如何实现？

- theodre7

Answer 1

使用高级索引：

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

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

或者使用numpy.put_along_axis：

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

np.put_along_axis(A, idx, 1, axis=1)

更新了 A：

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

- mozway