Accessing a numpy array with an array of indices-CodePudding

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.

CodePudding user response：

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]

CodePudding user response：

Adding on @hpaulj an alternative way is:

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