arr = np.arange(12).reshape((3, 2, 2)) indices = np.array([0, 1, 1]) expected_outcome = np.array([[0, 1], [6, 7], [10, 11]])
I’m trying to index this array of shape (3,2,2) with an array of shape (3) containing the y-index of the value I want to get. I tried to make it work with for in statement, but is there an elegant way to do it with numpy?
Advertisement
Answer
So you want arr[0,0,:], arr[1,1,:], arr[2,1,:]?
How about
In [179]: arr[[0,1,2], [0,1,1]]
Out[179]:
array([[ 0, 1],
[ 6, 7],
[10, 11]])