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]]
Asked By: voo

||

Answer #1:

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]])
Answered By: Divakar

Answer #2:

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.

Answered By: hpaulj

Answer #3:

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]]
Answered By: Ehsan

Answer #4:

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]]
Answered By: rnso
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 .

Leave a Reply

Your email address will not be published.