Home > Net >  Indexing a 3D np array with two 1D np array
Indexing a 3D np array with two 1D np array

Time:02-10

I have a 3D np.array of shape (i,j,k)and two 1D np.array contain the indexes in the j-axis and k-axis.

For example

idx = np.array([0, 1])
idy = np.array([1, 2])
array = np.arange(0, 18).reshape(2, 3, 3)
array([[[ 0,  1,  2],
        [ 3,  4,  5],
        [ 6,  7,  8]],

       [[ 9, 10, 11],
        [12, 13, 14],
        [15, 16, 17]]])

I want to have array[0,0,1], array[0,1,2], array[1,0,1] and array[1,1,2] which is np.array([1,5,10,14]).

CodePudding user response:

Is this enough:

In [197]: arr[:,idx,idy]
Out[197]: 
array([[ 1,  5],
       [10, 14]])
In [198]: arr[:,idx,idy].ravel()
Out[198]: array([ 1,  5, 10, 14])
  • Related