assigning values in a numpy array

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 2x3 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

Solution 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.

Solution 2

Use numpy.meshgrid() to make arrays of indexes that you can use to index into both your original array and the array of values for the third dimension.

import numpy as np
import scipy as sp
import scipy.stats.distributions

a = np.zeros((2,3,4))
z = sp.stats.distributions.randint.rvs(0, 4, size=(2,3))

xx, yy = np.meshgrid( np.arange(2), np.arange(3) )
a[ xx, yy, z[xx, yy] ] = 1
print a

I've renamed your array from x to a, and the array of indexes from y to z, for clarity.

EDIT: 4D example:

a = np.zeros((2,3,4,5))
z = sp.stats.distributions.randint.rvs(0, 4, size=(2,3))
w = sp.stats.distributions.randint.rvs(0, 5, size=(2,3))

xx, yy = np.meshgrid( np.arange(2), np.arange(3) )
a[ xx, yy, z[xx, yy], w[xx, yy] ] = 1

Comments

[user1857751]

Thanks for the response, but the problem is that I need to do this efficiently for large arrays. In other words, I want to avoid for loops.

[user1857751]

I need code that will handle any dimensions for x, i.e. any number of dimensions or any number of elements in each dimension.

[James Waldby - jwpat7]

I'm not highly familiar with numpy so I don't know if it has a relevant builtin function to do the job. Of course there are related routines argwhere, nonzero, where, and extract that can do the opposite eg can find the locations of non-zeros in an array

[user1857751]

This definitely works faster than the for loops I was using, and it seems to scale really well with the number of dimensions and number of elements in each dimension. Thanks.

[user1857751]

Thanks for the response. Do you know how to generalize this to a with any number of dimensions?

[Alex I]

@user1857751: yes, of course, it generalizes in a straightforward way. meshgrid to create vectors of indexes in all the dimensions that you do not want to lookup, and then use those to index the original array and the indexes you do want to look up.