Numpy: access to a two-dimensional array with a 2D array of indices

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 data matrix for accessing these indices

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

and I want to use indices a to get b entries. Simply:

 >>> b[a] 

doesn't work, as it gives one row from b for each entry in a , i.e.

 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 an index pair on the last axis a to give two indices b :

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

Is there a clean way to do this, or do I need to change the form of b and combine the columns of a accordingly?

In my real problem, a has about 5 million records, and b 100 per 100, I would like to avoid loops.

Answer 1

Actually, this works:

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

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

Answer 2

In this case, it works.

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

Answer 3

A more general solution, when you want to use a 2D array of form indices (n, m) with an arbitrary large dimension m , is called inds , so that to access the elements of another 2D array of form (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) 

You can verify this using, for example:

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