numpy установить значение с другим многомерным массивом в качестве индекса

Question

Предположим, что существует 4-х мерный массив idx1, сохраняет индекс 5-го измерения для другого 5-ти размерного массива нулей1. например:

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 расширенным индексированием.

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