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

Answers:

Thank you for visiting the Q&A section on Magenaut. Please note that all the answers may not help you solve the issue immediately. So please treat them as advisements. If you found the post helpful (or not), leave a comment & I’ll get back to you as soon as possible.

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

Method 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.

Method 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]]

Method 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]]

All methods was sourced from stackoverflow.com or stackexchange.com, is licensed under cc by-sa 2.5, cc by-sa 3.0 and cc by-sa 4.0