NumPy数组索引4D数组

Question

我有一个大小为（2,3,4,4）且填充了零的四维数组'a'。

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

我还有一个大小为(2,3,4)的三维数组'b'，其中包含一些索引值（都在0到3之间）。

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

我可以通过以下展示的3个循环实现这一点：

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

- Babak

Answer 1

是的，你可以以向量化的形式完成这个操作：

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

这应该更容易地推广到更高阶的ndarrays。

- Francisco Vargas