多维 numpy 数组中的索引

Question

提前感谢您的帮助！

我有一个 4D numpy 数组 a

a = np.random.random_integers(1,5,(2,3,3,4))
>>> a
array([[[[5, 2, 5, 3],
         [1, 3, 1, 2],
         [4, 3, 1, 2]],

        [[3, 1, 3, 5],
         [2, 5, 2, 3],
         [1, 4, 1, 3]],

        [[4, 2, 4, 3],
         [1, 3, 5, 2],
         [3, 5, 4, 4]]],


       [[[1, 1, 2, 3],
         [2, 1, 5, 3],
         [1, 5, 5, 4]],

        [[1, 2, 3, 3],
         [5, 2, 4, 3],
         [4, 3, 4, 5]],

        [[4, 3, 5, 3],
         [3, 5, 2, 4],
         [4, 3, 3, 1]]]])

我有一个 3D numpy 数组 b，其形状与 a 的最后三个维度相同，由 0 和 1 组成。

b = np.random.random_integers(0,1,(3,3,4))
>>> b
array([[[0, 0, 0, 0],
        [0, 0, 0, 1],
        [1, 0, 1, 0]],

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

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

对于a中的每个(3,3,4)数组，我想选择在数组b中相同位置具有值为1的元素，并根据另一个数组c更新a中的这些值，其形状为a 的第一个维度：

c = np.array([-1, -2])

因此，数组 b 中值为 1 的 a[0,:] 中的值将变为 -1 (c[0])

以及 a[1,:] 中的值。数组 b 中值为 1 的值将变为 -2 (c[1])

最终我想以有效的方式获得一个与 a 形状相同且具有更新值的数组

>>> return 
array([[[[5, 2, 5, 3],
         [1, 3, 1, -1],
         [-1, 3, -1, 2]],

        [[3, -1, -1, 5],
         [-1, -1, 2, -1],
         [-1, -1, 1, -1]],

        [[4, 2, 4, -1],
         [1, -1, -1, 2],
         [-1, 5, -1, 4]]],


       [[[1, 1, 2, 3],
         [2, 1, 5, -2],
         [-2, 5, -2, 4]],

        [[1, -2, -2, 3],
         [-2, -2, 4, -2],
         [-2, -2, 4, -2]],

        [[4, 3, 5, -2],
         [3, -2, -2, 4],
         [-2, 3, -2, 1]]]])

Answer 1

简单来说：

for i in range(a.shape[0]):
    a[i][np.where(b)] = c[i]

print(a)

---

[[[[ 5  2  5  3]
   [ 1  3  1 -1]
   [-1  3 -1  2]]

  [[ 3 -1 -1  5]
   [-1 -1  2 -1]
   [-1 -1  1 -1]]

  [[ 4  2  4 -1]
   [ 1 -1 -1  2]
   [-1  5 -1  4]]]


 [[[ 1  1  2  3]
   [ 2  1  5 -2]
   [-2  5 -2  4]]

  [[ 1 -2 -2  3]
   [-2 -2  4 -2]
   [-2 -2  4 -2]]

  [[ 4  3  5 -2]
   [ 3 -2 -2  4]
   [-2  3 -2  1]]]]

Answer 2

使用广播和

numpy.where

:

tmp = b[None] * c[:,None,None,None]
out = np.where(tmp, tmp, a)

输出：

array([[[[ 5,  2,  5,  3],
         [ 1,  3,  1, -1],
         [-1,  3, -1,  2]],

        [[ 3, -1, -1,  5],
         [-1, -1,  2, -1],
         [-1, -1,  1, -1]],

        [[ 4,  2,  4, -1],
         [ 1, -1, -1,  2],
         [-1,  5, -1,  4]]],


       [[[ 1,  1,  2,  3],
         [ 2,  1,  5, -2],
         [-2,  5, -2,  4]],

        [[ 1, -2, -2,  3],
         [-2, -2,  4, -2],
         [-2, -2,  4, -2]],

        [[ 4,  3,  5, -2],
         [ 3, -2, -2,  4],
         [-2,  3, -2,  1]]]])