Select one dimension of Multidimensional array with list - numpy

Question

I have a 3D array of shape (800,5,4) like:

arr = array([[35.  , 33.  , 33.  ,  0.15],
       [47.  , 47.  , 44.  ,  0.19],
       [49.  , 56.  , 60.  ,  0.31],
       ...,
       [30.  , 27.  , 25.  ,  0.07],
       [54.  , 49.  , 42.  ,  0.14],
       [33.  , 30.  , 28.  ,  0.22]])

I have a 1D array of indeces for the second dimension (so they range from 0 to 4) like this:

indeces = [0,3,2,0,1,1,1,0,...,0,1,2,2,4,3]

I want to select the idx item from the second dimension, and get back an array of shape (800,4)

I have tried the following but could not make it work:

indexed = arr[:,indeces,:]

What am I missing?

Answer 1

In [178]: arr = np.arange(24).reshape(2,3,4)    

If I have a list of 7 items:

In [179]: idx = [0,1,1,2,2,0,1]                                                                  
In [180]: arr[:,idx,:]                                                                           
Out[180]: 
array([[[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11],
        [ 8,  9, 10, 11],
        [ 0,  1,  2,  3],
        [ 4,  5,  6,  7]],

       [[12, 13, 14, 15],
        [16, 17, 18, 19],
        [16, 17, 18, 19],
        [20, 21, 22, 23],
        [20, 21, 22, 23],
        [12, 13, 14, 15],
        [16, 17, 18, 19]]])
In [181]: _.shape                                                                                
Out[181]: (2, 7, 4)

To produce a (2,4) result, we have to pick one element on the 2nd dim for each of pair from the other dimensions.

A general case would be to make idx a (2,4) array, and index with dimensions that also broadcast to (2,4):

In [182]: idx = np.array([0,1,1,2,2,0,1,0]).reshape(2,4)                                         
In [183]: arr[np.arange(2)[:,None],idx,np.arange(4)]                                             
Out[183]: 
array([[ 0,  5,  6, 11],
       [20, 13, 18, 15]])
In [184]: _.shape                                                                                
Out[184]: (2, 4)

Or we could pick with a scalar:

In [185]: arr[:,2,:]                                                                             
Out[185]: 
array([[ 8,  9, 10, 11],
       [20, 21, 22, 23]])

@a_guest showed how to do this with an idx that matches the 1st dimension (and slices the last).

One way or other your idx has to map or broadcast with the other dimensions.

Answer 2

You need to also select the remaining dimensions "row-by-row". That is you need to supply an index array of the following form: range(arr.shape[0]) . For example:

>>> a = np.arange(27).reshape(3, 3, 3)
>>> i = [0, 1, 2]
>>> a[range(a.shape[0]), i, :]
array([[ 0,  1,  2],
       [12, 13, 14],
       [24, 25, 26]])
>>> a[range(a.shape[0]), i, range(a.shape[2])]
array([ 0, 13, 26])