CODEWITHSUNDEEP
pythonarraysnumpy
Lucas
Lucas

Reputation: 1948

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([[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, 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 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.

Upvotes: 2

Views: 1221

Answers (3)

Ataxias
Ataxias

Reputation: 1193

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

Upvotes: 0

colcarroll
colcarroll

Reputation: 3682

Actually, this works:

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

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

Upvotes: 2

M4rtini
M4rtini

Reputation: 13539

For this case, this works

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

Upvotes: 1

Related Questions