Indexing One Array by Another in Numpy

Indexing one array by another in numpy

EDIT: np.take_along_axis is a builtin function for this use case implemented since numpy 1.15. See @hpaulj 's answer below for how to use it.

---

You can use NumPy's advanced indexing -

A[np.arange(A.shape[0])[:,None],B]

One can also use linear indexing -

m,n = A.shape
out = np.take(A,B + n*np.arange(m)[:,None])

Sample run -

In [40]: A
Out[40]: 
array([[2, 4, 5, 3],
       [1, 6, 8, 9],
       [8, 7, 0, 2]])

In [41]: B
Out[41]: 
array([[0, 0, 1, 2],
       [0, 3, 2, 1],
       [3, 2, 1, 0]])

In [42]: A[np.arange(A.shape[0])[:,None],B]
Out[42]: 
array([[2, 2, 4, 5],
       [1, 9, 8, 6],
       [2, 0, 7, 8]])

In [43]: m,n = A.shape

In [44]: np.take(A,B + n*np.arange(m)[:,None])
Out[44]: 
array([[2, 2, 4, 5],
       [1, 9, 8, 6],
       [2, 0, 7, 8]])

Indexing a numpy array with another array

Try using np.arange(a.shape[0]) as row index, with the array b as column index:

>>> a[np.arange(a.shape[0]), b].shape
(10,)

Get indices of element of one array using indices in another array

I have a solution similar to that of Andras based on np.argmax and np.arange. Instead of "indexing the index" I propose to add a piecewise offset to the result of np.argmax:

import numpy as np
a = np.array([[[7, 9],
               [19, 18]],
              [[24, 5],
               [18, 11]]])
off = np.arange(0, a.size, a.shape[2]).reshape(a.shape[0], a.shape[1])

>>> off
array([[0, 2],
       [4, 6]])

This results in:

>>> a.argmax(-1) + off
array([[1, 2],
       [4, 6]])

Or as a one-liner:

>>> a.argmax(-1) + np.arange(0, a.size, a.shape[2]).reshape(a.shape[0], a.shape[1])
array([[1, 2],
       [4, 6]])

Indexing multidimensional Numpy array with another array

Tim Roberts summarized it excellently in the comments, which I'll quote here in case the comments are cleaned up:

It's subtle. The second example is concatenating a set of slices. If you print a[0,:,0], you'll see a 3-element slice, equal to the first row of the final b. Same with a[0,:,1]. The magic indexing takes those 5 3-element slices and returns then in a new array. A set of 5 slices is different from the array subset in a[0,:,0:5].

In addition, if you notice, the two different indexing methods actually produce equivalent results; they're just transposed versions of each other. So a[0, :, np.arange(5)] == a[0, :, 0:5].T and a[0, :, np.arange(5)].T == a[0, :, 0:5].

Index a numpy array with another array

Just use a tuple:

>>> A[(3, 1)]
8
>>> A[tuple(ind)]
8

The A[] actually calls the special method __getitem__:

>>> A.__getitem__((3, 1))
8

and using a comma creates a tuple:

>>> 3, 1
(3, 1)

Putting these two basic Python principles together solves your problem.

You can store your index in a tuple in the first place, if you don't need NumPy array features for it.

Indexing numpy array with another numpy array

According the NumPy tutorial, the correct way to do it is:

a[tuple(b)]

Numpy: For every element in one array, find the index in another array

As Joe Kington said, searchsorted() can search element very quickly. To deal with elements that are not in x, you can check the searched result with original y, and create a masked array:

import numpy as np
x = np.array([3,5,7,1,9,8,6,6])
y = np.array([2,1,5,10,100,6])

index = np.argsort(x)
sorted_x = x[index]
sorted_index = np.searchsorted(sorted_x, y)

yindex = np.take(index, sorted_index, mode="clip")
mask = x[yindex] != y

result = np.ma.array(yindex, mask=mask)
print result

the result is:

[-- 3 1 -- -- 6]

How to index a numpy array with another numpy array in python

A widely used approach for this kind of problem is to construct a boolean mask, comparing the index array with the appropriate arange:

In [619]: mask = np.arange(4)[:,None]>=d
In [620]: mask
Out[620]: 
array([[False, False, False],
       [False,  True, False],
       [ True,  True, False],
       [ True,  True,  True]])
In [621]: a[mask]
Out[621]: array([ 5,  7,  8, 10, 11, 12])
In [622]: a[mask] = 0
In [623]: a
Out[623]: 
array([[1, 2, 3],
       [4, 0, 6],
       [0, 0, 9],
       [0, 0, 0]])

That's not necessarily faster than a row (or in this case column) iteration. Since slicing is basic indexing, it may be faster, even if done several times.

In [624]: for i,v in enumerate(d):
     ...:     print(a[v:,i])
     ...: 
[0 0]
[0 0 0]
[0]

Generally if a result involves multiple arrays or lists with different lengths, there isn't a "neat" multidimensional solution. Either iterate over those lists, or step back and "think outside the box".

---

Related Topics