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