Numpy: 2d Array Access With 2d Array Of Indices

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(

Solution 1:

Actually, this works:

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

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

Solution 2:

For this case, this works

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

Solution 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 indsoffset = 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 CResult = 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, e.g.:

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