CODEWITHSUNDEEP
pythonnumpymultidimensional-arraynumpy-ndarray
illpack
illpack

Reputation: 127

Select one dimension of Multidimensional array with list - numpy

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?

Upvotes: 1

Views: 7285

Answers (2)

hpaulj
hpaulj

Reputation: 231335

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.

Upvotes: 1

a_guest
a_guest

Reputation: 36239

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])

Upvotes: 0

Related Questions