NumPy array indexing 4D array

Question

I have a 4D array "a" of size (2,3,4,4) filled with zeros.

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

I also have a 3D array "b" of size (2,3,4) that carries some index values ​​(all between 0 and 3).

What I want to do is replace the element of each last array with 'a' (4-dimensional dimension 'a'), which corresponds to the index in 'b', with 1.

I can do this with 3 for loops, as shown below:

 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 

But I was wondering if there was any way to avoid the loop at all. Something like:

 a[b] = 1 

Answer 1

Yes, you can do this in vector form:

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

This should more easily generalize higher order ndarrays.