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?

Your result will be of size (800,len(indeces), 4) as you are selecting from that middle index many times. If you want it to be (800,4) than you can only select one index like arr[:,indeces[0],:]. — jhso
– jhso, Commented Jan 20, 2020 at 22:39

hpaulj · Accepted Answer · 2020-01-20 23:54:39Z

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.

Would you consider stackoverflow.com/questions/68327803/… a duplicate of this question?

a_guest · Accepted Answer · 2020-01-20 23:15:02Z

0

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

answered Jan 20, 2020 at 23:15

a_guest

36.7k15 gold badges75 silver badges137 bronze badges

Collectives™ on Stack Overflow

Select one dimension of Multidimensional array with list - numpy

2 Answers 2

1 Comment

Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

1 Comment

Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related