Is it possible to indexing 2d array by another 2d array row by row without loop in NumPy?

Viewed 89

Let's say I have NumPy array with indexes:

array([[1, 3],
       [2, 5]])

Also, I have NumPy array with zeros:

array([[0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0.]])

I need to take the first row from the array with indexes and set 1 for these indexes in the first row of the array with zeros. After this first row must look like [0, 1, 0, 1, 0, 0].

Then I need to do also for the second row in the array with zeros and use the second row from the array with indexes respectively. After this second row must look like [0, 0, 1, 0, 0, 1].

So, I want to get array like this:

array([[0, 1, 0, 1, 0, 0],
       [0, 0, 1, 0, 0, 1]])

How can I do this without any loop?

2 Answers

Given the array of indices i and the target array x.

>>> i = np.array([[1, 3],
                  [2, 5]])

>>> x = np.array([[0., 0., 0., 0., 0., 0.],
                  [0., 0., 0., 0., 0., 0.]])

You could use np.put_along_axis:

numpy.put_along_axis(arr, indices, values, axis)
Put values into the destination array by matching 1d index and data slices.

>>> np.put_along_axis(x, i, values=1, axis=1)

>>> x
array([[0., 1., 0., 1., 0., 0.],
       [0., 0., 1., 0., 0., 1.]])

Prior to the addition of np.put_along_axis you could have used an arange index on the first dimension:

In [542]: x=np.array([[1, 3],[2, 5]]); y = np.zeros((2,6),int)

This indexes a block:

In [543]: y[np.arange(2),x]
Out[543]: 
array([[0, 0],
       [0, 0]])

or using it to set:

In [544]: y[np.arange(2),x]=1
In [545]: y
Out[545]: 
array([[0, 1, 1, 0, 0, 0],
       [0, 0, 0, 1, 0, 1]])

Oops, that's not quite what you want. Let's try transposing x first:

In [546]: x=np.array([[1, 3],[2, 5]]); y = np.zeros((2,6),int)
In [547]: y[np.arange(2),x.T]=1
In [548]: y
Out[548]: 
array([[0, 1, 0, 1, 0, 0],
       [0, 0, 1, 0, 0, 1]])
Related