# Indexing one array by another in numpy

Posted on

Solving problem is about exposing yourself to as many situations as possible like Indexing one array by another in numpy and practice these strategies over and over. With time, it becomes second nature and a natural way you approach any problems in general. Big or small, always start with a plan, use other strategies mentioned here till you are confident and ready to code the solution.
In this post, my aim is to share an overview the topic about Indexing one array by another in numpy, which can be followed any time. Take easy to follow this discuss.

Indexing one array by another in numpy

Suppose I have a matrix A with some arbitrary values:

``````array([[ 2, 4, 5, 3],
[ 1, 6, 8, 9],
[ 8, 7, 0, 2]])
``````

And a matrix B which contains indices of elements in A:

``````array([[0, 0, 1, 2],
[0, 3, 2, 1],
[3, 2, 1, 0]])
``````

How do I select values from A pointed by B, i.e.:

``````A[B] = [[2, 2, 4, 5],
[1, 9, 8, 6],
[2, 0, 7, 8]]
``````

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]])
``````

More recent versions have added a `take_along_axis` function that does the job:

``````A = np.array([[ 2, 4, 5, 3],
[ 1, 6, 8, 9],
[ 8, 7, 0, 2]])
B = np.array([[0, 0, 1, 2],
[0, 3, 2, 1],
[3, 2, 1, 0]])
np.take_along_axis(A, B, 1)
Out[]:
array([[2, 2, 4, 5],
[1, 9, 8, 6],
[2, 0, 7, 8]])
``````

There’s also a `put_along_axis`.

I know this is an old question, but another way of doing it using indices is:

``````A[np.indices(B.shape)[0], B]
``````

output:

``````[[2 2 4 5]
[1 9 8 6]
[2 0 7 8]]
``````

Following is the solution using `for` loop:

``````outlist = []
for i in range(len(B)):
lst = []
for j in range(len(B[i])):
lst.append(A[i][B[i][j]])
outlist.append(lst)
outarray = np.asarray(outlist)
print(outarray)
``````

Above can also be written in more succinct list comprehension form:

``````outlist = [ [A[i][B[i][j]] for j in range(len(B[i]))]
for i in range(len(B))  ]
outarray = np.asarray(outlist)
print(outarray)
``````

Output:

``````[[2 2 4 5]
[1 9 8 6]
[2 0 7 8]]
``````
The answers/resolutions are collected from stackoverflow, are licensed under cc by-sa 2.5 , cc by-sa 3.0 and cc by-sa 4.0 .