Numpy：使用二维索引数组访问二维数组

Question

I have two arrays, one is a matrix of index pairs,

a = array([[[0,0],[1,1]],[[2,0],[2,1]]], dtype=int)

and another which is a matrix of data to access at these indices

b = array([[1,2,3],[4,5,6],[7,8,9]])

and I want to able to use the indices of a to get the entries of b . Just doing:

>>> b[a]

does not work, as it gives one row of b for each entry in a , ie

array([[[[1,2,3],
         [1,2,3]],

        [[4,5,6],
         [4,5,6]]],


       [[[7,8,9],
         [1,2,3]],

        [[7,8,9],
         [4,5,6]]]])

when I would like to use the index pair in the last axis of a to give the two indices of b :

array([[1,5],[7,8]])

Is there a clean way of doing this, or do I need to reshape b and combine the columns of a in a corresponding manner?

In my actual problem a has about 5 million entries, and b is 100-by-100, I'd like to avoid for loops.

Answer 1

Actually, this works:

b[a[:, :, 0],a[:, :, 1]]

Gives array([[1, 5], [7, 8]]) .

Answer 2

For this case, this works

tmp =  a.reshape(-1,2)
b[tmp[:,0], tmp[:,1]] 

Answer 3

A more general solution, whenever you want to use a 2D array of indices of shape (n,m) with arbitrary large dimension m , named inds , in order to access elements of another 2D array of shape (n,k), named B :

# array of index offsets to be added to each row of inds
offset = np.arange(0, inds.size, inds.shape[1])

# numpy.take(B, C) "flattens" arrays B and C and selects elements from B based on indices in C
Result = np.take(B, offset[:,np.newaxis]+inds)

Another solution, which doesn't use np.take and I find more intuitive, is the following:

B[np.expand_dims(np.arange(B.shape[0]), -1), inds]

The advantage of this syntax is that it can be used both for reading elements from B based on inds (like np.take ), as well as for assignment.

You can test this by using, eg:

B = 1/(np.arange(n*m).reshape(n,-1) + 1)
inds = np.random.randint(0,B.shape[1],(B.shape[0],B.shape[1]))