使用 2D 索引数组对 3D 数组值进行 Numpy 赋值答案

【问题标题】：Numpy assignment of 3D array values using 2D index array使用 2D 索引数组对 3D 数组值进行 Numpy 赋值
【发布时间】：2021-01-29 12:34:49
【问题描述】：

我有一个 3D numpy 零数组，尺寸为 CxHxW（在本例中，C=4、H=2 和 W=3）：

A = np.array([[[0, 0, 0],
               [0, 0, 0]],
              [[0, 0, 0],
               [0, 0, 0]]
              [[0, 0, 0],
               [0, 0, 0]]
              [[0, 0, 0],
               [0, 0, 0]]])

我还有一个二维索引数组，尺寸为 HxW，因此数组中的每个值都是 [0, C-1] 之间的有效索引

B = np.array([[2, 3, 0], 
              [3, 1, 2]])

有没有一种快速的方法，使用矢量化来修改数组 A，使得 A[B[i][j]][i][j] = 1，对于所有有效的 i，j？

A = np.array([[[0, 0, 1],
               [0, 0, 0]],
              [[0, 0, 0],
               [0, 1, 0]]
              [[1, 0, 0],
               [0, 0, 1]]
              [[0, 1, 0],
               [1, 0, 0]]])

谢谢！

【问题讨论】：

标签： python numpy indexing vectorization

【解决方案1】：

你好像在找put_along_axis:

np.put_along_axis(A, B[None,...], 1, 0)

请注意，第二个参数必须与第一个参数具有相同的维数，这就是为什么使用B[None,...] 而不是B。

【讨论】：