NumPy array indexing a 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 every last array in 'a' (the 4th dimension of 'a') that 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 is anyway I can avoid looping at all. Something similar to:

a[b] = 1

Solution

Yes you can do this in a vectorized form:

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

This should generalize more easily to higher order ndarrays.