How to index an N dimensional array using an N-1 dimensional array of indices-CodePudding

I have a 4-dimensional array rand that I would like to index along the 3rd dimension using a 3-dimensional array idx. In the example below, the output should be of shape [32,100,3]

import numpy as np

# values to be indexed
array = np.random.rand(32, 100, 4, 3)

# values to be indexed along the 3rd dimension of array
idx = np.random.randint(low=0, high=3, size=(32,100,1))

indexed_array = array[idx] # Output here should be [32, 100, 3]

What am I doing wrong? A solution would be to loop over every dimension but is there a better alternative?

CodePudding user response：

In [61]: array = np.random.rand(32, 100, 4, 3)
    ...: 
    ...: # values to be indexed along the 3rd dimension of array
    ...: idx = np.random.randint(low=0, high=3, size=(32, 100, 1))
In [62]: array.shape
Out[62]: (32, 100, 4, 3)
In [63]: idx.shape
Out[63]: (32, 100, 1)

With:

In [64]: array[idx].shape
Out[64]: (32, 100, 1, 100, 4, 3)

we are applying the array idx to the first dimension of array.

idx is indices for the 3rd dimension, not the first. To apply it we need to specify compatible arrays for the first 2 dimensions (in the broadcasting sense):

In [66]: array[np.arange(32)[:, None,None], np.arange(100)[:,None], idx].shape
Out[66]: (32, 100, 1, 3)
In [67]: 
In [67]: array[np.arange(32)[:, None], np.arange(100), idx[:,:,0]].shape
Out[67]: (32, 100, 3)

There is a newish function that may streamline this (for some people):

In [70]: np.take_along_axis(array, idx[:, :, :, None], axis=2).shape
Out[70]: (32, 100, 1, 3)

When indexing an array, there's an important distinction between a tuple and an array (or list).

arr[1,2,3]  # index each of the 3 dimensions
arr[(1,2,3)]  # equivalent
arr[[1,2,3]]   # index just the first dimension with a list