如何根据另一个二维数组中给出的索引对二维数组进行切片

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?

Answer 1

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]

Answer 2

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