how to slice a 2d array based on indices given in another 2d array

Question

I have an MxN array called A which stores the data I want. I have another M x N2 array B which stores array indices, and N2<N. Each row of B stores the indices of the elements I want to get from A for that row. For example, the following code works for me:

A_reduced = np.zeros((M,N2))
for i in range(M):
    A_reduced[i,:] = A[i,B[i,:]]

Are there any 'vectorized' ways to extract the desired elements from A based on B instead of looping through each row?

Just clarifying. let's say A_reduced is `[0, 0, 0]` and B is `[1, 2]`. Would you expect A_reduced to now be `[0, 0]` — zerecees, Apr 24 '20 at 03:48
Why do you want to vectorize it? Do you think it'll make the code run faster or is there a different reason? If you want to make it easy to read you can do something like `A_reduced[i,:] = [A[i,B[i,:]] for i in range(M)]`. I think numpy has a vectorize class but that might overcomplicate: https://docs.scipy.org/doc/numpy/reference/generated/numpy.vectorize.html — Ariel A, Apr 24 '20 at 03:48

Quang Hoang · Answer 1 · 2020-04-24T04:08:10.110

You can exploit array indexing and use reshape:

# set up M=N=4, N2=2
a = np.arange(16).reshape(4,4)
b = np.array([[1,2],[0,1],[2,3],[1,3]])

row_idx = np.repeat(np.arange(b.shape[0]),b.shape[1])
col_idx = b.ravel()

# output:
a[row_idx, col_idx].reshape(b.shape)

Output:

array([[ 1,  2],
       [ 4,  5],
       [10, 11],
       [13, 15]])

Update: Another similar solution

row_idx = np.repeat(np.arange(b.shape[0]),b.shape[1]).reshape(b.shape)

# output
a[row_idx,b]

score 1 · Answer 2 · answered Apr 24 '20 at 05:41

In [203]: A = np.arange(12).reshape(3,4)                                                               
In [204]: B = np.array([[0,2],[1,3],[3,0]])

Your row iteration:

In [207]: A_reduced = np.zeros((3,2),int)                                                              
In [208]: for i in range(3): 
     ...:     A_reduced[i,:] = A[i, B[i,:]] 
     ...:                                                                                              
In [209]: A_reduced                                                                                    
Out[209]: 
array([[ 0,  2],
       [ 5,  7],
       [11,  8]])

A 'vectorized' version:

In [210]: A[np.arange(3)[:,None], B]                                                                   
Out[210]: 
array([[ 0,  2],
       [ 5,  7],
       [11,  8]])

and streamlined with a newish function:

In [212]: np.take_along_axis(A,B,axis=1)                                                               
Out[212]: 
array([[ 0,  2],
       [ 5,  7],
       [11,  8]])

how to slice a 2d array based on indices given in another 2d array

2 Answers2