NumPy数组索引4D数组

Question

我有一个4D数组'a'大小(2,3,4,4)填充零.

import numpy as np
a = np.zeros((2,3,4,4))

我还有一个大小为(2,3,4)的3D数组'b',它带有一些索引值(全部在0到3之间).

我想要做的是将'a'中的每个最后一个数组的元素('a'的第四个维度)替换为'b'中的索引,使用1.

我可以使用3 for循环执行此操作,如下所示:

for i in a.shape[0]:
    for j in a.shape[1]:
        for z in a.shape[2]:
            a[i,j,z][b[i,j,z]] = 1 

但我想知道无论如何我都可以避免循环.类似的东西:

a[b] = 1

Answer 1

是的，您可以以矢量形式进行此操作：

p,m,n,r = a.shape
a.reshape(-1,r)[np.arange(p*m*n),b.ravel()] = 1

这应该更容易推广到高阶ndarray。