Indexing 3d numpy array with 2d array Indexing 3d numpy array with 2d array numpy numpy

Indexing 3d numpy array with 2d array


The shape of the result from fancy index and broadcasting is the shape of the indexing array. You need passing 2d array for each axis of arr_3d

ax_0 = np.arange(arr_3d.shape[0])[:,None]ax_1 = np.arange(arr_3d.shape[1])[None,:]arr_3d[ax_0, ax_1, arr_2d]Out[1127]:array([[ 3,  6,  8],       [14, 19, 22]])


In [107]: arr_3d = np.arange(2*3*4).reshape(2,3,4)                                                           In [108]: arr_2d = np.array(([3,2,0], [2,3,2]))                                                              In [109]: arr_2d.shape                                                                                       Out[109]: (2, 3)In [110]: arr_3d[[[0],[1]],[0,1,2],arr_2d]                                                                   Out[110]: array([[ 3,  6,  8],       [14, 19, 22]])

[[0],[1]],[0,1,2] broadcast with each other to index a (2,3) block, the same size as `arr_2d.

ix_ can be used to construct those 2 indices:

In [114]: I,J = np.ix_(range(2), range(3))                                                                   In [115]: I,J                                                                                                Out[115]: (array([[0],        [1]]), array([[0, 1, 2]]))In [116]: arr_3d[I, J, arr_2d]                                                                               Out[116]: array([[ 3,  6,  8],       [14, 19, 22]])