索引 4D 数组的 NumPy 数组

Question

我有一个大小为 (2,3,4,4) 的 4D 数组“a”，其中填充了零。

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

我还有一个大小为 (2,3,4) 的 3D 数组“b”，它带有一些索引值（都在 0 和 3 之间）。

我想要做的是用 1 替换 'a' 中每个最后一个数组的元素（'a' 的第 4 维），它对应于 'b' 中的索引。

我可以用 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。