numpy set value with another multiple dimension array as index

user_1248918 0
Viewed 62

Assume there is a 4 dimension array idx1, stores 5 th dimension index for another 5 dimension array zeros1. like:

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

I want to realize

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

How can I do it with numpy advanced indexing.

numpy python
1 Answers

Use open-range arrays and index to assign -

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

Alternatively, in one-line -

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