numpyvectorizationnumpy-ndarrayarray-broadcastingnumpy-indexing

A faster way than "for loop" - Indexing 3d numpy array by 2d array

I have a 3d (3, 2, 3) array and first dimension(3) is flexible it can be any size.

arr = np.array(
[[[56, 24, 32],
  [56, 24, 32]],
 [[51, 27, 72],
  [51, 27, 72]],
 [[36, 14, 49],
  [36, 14, 49]]])

Indexing array is (2,3):

idxs = np.array(
[[1, 0, 2],
 [2, 1, 0]])

I wanna index arr by idxs. Expected result is :

[[[24 56 32]
  [32 24 56]]
 [[27 51 72]
  [72 27 51]]
 [[14 36 49]
  [49 14 36]]])

If i use for loop like below it takes a lot of time :

for i, arr2d in enumerate(arr):
    for j, (arr1d, idx) in enumerate(zip(arr2d, idxs)):
        arr[i, j] = arr1d[idx]

So my question is : How can i speed up this process ?

Solution

  • Use np.take_along_axis

    np.take_along_axis(arr, idxs[None, ...], 2)
Out[]: 
array([[[24, 56, 32],
        [32, 24, 56]],

       [[27, 51, 72],
        [72, 27, 51]],

       [[14, 36, 49],
        [49, 14, 36]]])