Question

I have a numpy array of zeros. For concreteness, suppose it’s 2x3x4:

x = np.zeros((2,3,4))

and suppose I have a 2×3 array of random integers from 0 to 3 (the index of the 3rd dimension of x).

>>> y = sp.stats.distributions.randint.rvs(0, 4, size=(2,3))
>>> y
[[2 1 0]
 [3 2 0]]

How do I do the following assignments efficiently (edit: something that doesn’t use for loops and works for x with any number of dimensions and any number of elements in each dimension)?

>>> x[0,0,y[0,0]]=1
>>> x[0,1,y[0,1]]=1
>>> x[0,2,y[0,2]]=1
>>> x[1,0,y[1,0]]=1
>>> x[1,1,y[1,1]]=1
>>> x[1,2,y[1,2]]=1
>>> x
array([[[ 0.,  0.,  1.,  0.],
        [ 0.,  1.,  0.,  0.],
        [ 1.,  0.,  0.,  0.]],

       [[ 0.,  0.,  0.,  1.],
        [ 0.,  0.,  1.,  0.],
        [ 1.,  0.,  0.,  0.]]])

Thanks,
James

Answer 1

At the moment, I can only think of the “simple” version, which involves flattening along the first two dimensions. This code should work:

shape_last = x.shape[-1]
x.reshape((-1, shape_last))[np.arange(y.size), y.flatten()] = 1

This yields (with my randomly-generated y):

array([[[ 0.,  0.,  0.,  1.],
        [ 0.,  0.,  1.,  0.],
        [ 0.,  1.,  0.,  0.]],

       [[ 0.,  1.,  0.,  0.],
        [ 0.,  0.,  0.,  1.],
        [ 0.,  1.,  0.,  0.]]])

The key is, if you do an indexing using multiple numpy arrays (advanced indexing), numpy will use pairs of indices to index into the array.

Of course, make sure x and y are both either C-order or F-order — otherwise, the calls to reshape and flatten might give different orders.