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

illya · Accepted Answer · 2012-11-27 22:14:39Z

2

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.

answered Nov 27, 2012 at 22:14

illya

7755 silver badges8 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

user1857751 Over a year ago

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.

Alex I · Accepted Answer · 2012-11-28 00:23:10Z

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

edited Nov 28, 2012 at 0:23

answered Nov 27, 2012 at 22:33

Alex I

20.4k12 gold badges100 silver badges166 bronze badges

2 Comments

user1857751 Over a year ago

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

Alex I Over a year ago

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

James Waldby - jwpat7 · Accepted Answer · 2012-11-27 22:01:40Z

0

x = np.zeros((2,3,4))
y=np.array([[2, 1, 0],[3, 2, 0]]) # or y=sp.stats...
for i in range(2):
    for j in range(3):
        x[i,j,y[i,j]]=1

will produce the desired result, IIRC. If the array dimensions never change, consider replacing the two for loops and their burden by

for j in range(3):
    x[0,j,y[0,j]] = x[1,j,y[1,j]] = 1

edited Nov 27, 2012 at 22:01

answered Nov 27, 2012 at 21:56

James Waldby - jwpat7

8,7312 gold badges24 silver badges38 bronze badges

3 Comments

user1857751 Over a year ago

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 Over a year ago

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 Over a year ago

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

Collectives™ on Stack Overflow

assigning values in a numpy array

3 Answers 3

1 Comment

2 Comments

3 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

1 Comment

2 Comments

3 Comments

Your Answer

Sign up or log in

Post as a guest