I am looking for a vectorized way to index a numpy.array
by numpy.array
of indices.
For example:
import numpy as np
a = np.array([[0,3,4],
[5,6,0],
[0,1,9]])
inds = np.array([[0,1],
[1,2],
[0,2]])
I want to build a new array, such that every row(i) in that array is a row(i) of array a
, indexed by row of array inds(i). My desired output is:
array([[ 0., 3.], # a[0][:,[0,1]]
[ 6., 0.], # a[1][:,[1,2]]
[ 0., 9.]]) # a[2][:,[0,2]]
I can achieve this with a loop:
def loop_way(my_array, my_indices):
new_array = np.empty(my_indices.shape)
for i in xrange(len(my_indices)):
new_array[i, :] = my_array[i][:, my_indices[i]]
return new_array
But I am looking for a pure vectorized solution.
When using arrays of indices to index another array, the shape of each index array should match the shape of the output array. You want the column indices to match inds
, and you want the row indices to match the row of the output, something like:
array([[0, 0],
[1, 1],
[2, 2]])
You can just use a single column of the above, due to broadcasting, so you can use np.arange(3)[:,None]
is the vertical arange
because None
inserts a new axis:
>>> np.arange(3)[:, None]
array([[0],
[1],
[2]])
Finally, together:
>>> a[np.arange(3)[:,None], inds]
array([[0, 3], # a[0,[0,1]]
[6, 0], # a[1,[1,2]]
[0, 9]]) # a[2,[0,2]]
It’s possible, although somewhat non-obvious to do this as follows:
>>> a[np.arange(a.shape[0])[:, None], inds]
array([[0, 3],
[6, 0],
[0, 9]])
The index np.arange(a.shape[0])
simply indexes the rows to which the array of column indices inds
applies. Appending [:, None]
modifies the shape of this array such that its shape is (a.shape[0], 1)
, i.e. each row index is in a separate row of a 1-column-wide 2D array.
The basic principle is that the number of dimensions in the index arrays must agree, and their shapes must also do so. See documentation for np.ix_
to get a feel for this.
If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!
Donate Us With