CODEWITHSUNDEEP
pythonnumpy
Daedeloth
Daedeloth

Reputation: 13

Accessing a numpy array with an array of indices

I'm trying to access a numpy array A using another array B providing the indices at each position:

A = np.array([[1,2],[3,4]])
B = np.array([[[0,0],[0,0]],[[0,1],[0,1]]])

Desired output:

C = array([[1,1],[3,3]])

I haven't gotten it to work using np.take() or the advanced indexing. I could do it iteratively but my arrays are on the order of 10**7 so I was hoping for a faster way.

Upvotes: 1

Views: 887

Answers (2)

Learning is a mess
Learning is a mess

Reputation: 8277

Adding on @hpaulj an alternative way is:

idx =  tuple(B[:,:,[1,0]].transpose(2,0,1))
A[idx]
# array([[1, 1], [3, 3]])

Upvotes: 0

hpaulj
hpaulj

Reputation: 231335

I probably should have insisted on seeing the iterative solution first, but here's the array one:

In [45]: A[B[:,:,1], B[:,:,0]]
Out[45]: 
array([[1, 1],
       [3, 3]])

I first tried A[B[:,:,0], B[:,:,1]], the natural order of the inner dimension. Your own code could have saved me that trial.

The key with advanced indexing is that you have to create or define separate arrays (broadcastable) for each dimension. We can think of that index as a tuple:

 idx = (B[:,:,0], B[:,:,1])
 A[idx]

Upvotes: 2

Related Questions