用另一个多维数组作为索引设置NumPy数组的值

Question

假设有一个四维数组idx1，存储另一个五维数组zeros1的5个维度索引。

N,T,H,W = idx1.shape
zeros1 = np.zeros( (N,T,H,W, 256) )
# it is guaranteed that idx1's value <256

我想实现

for n in range(N):
    for t in range(T):
        for h in range(H):
            for w in range(W):
                x = idx1[ n,t,h,w ]
                zeros1[n,t,h,w,x] = 1

我该如何使用numpy高级索引进行操作。

- worldterminator

Answer 1

使用开放范围数组和索引进行赋值 -

out = np.zeros( (N,T,H,W, 256) )
i,j,k,l = np.ogrid[:N,:T,:H,:W]
out[i,j,k,l,idx1] = 1

或者，简单来说：

out[tuple((np.ogrid[:N,:T,:H,:W]+[idx1]))] = 1

- Divakar