Get columns of 2D ndarray defined by index' in another 2D ndarray

Question

I have an ndarray arr = np.array([[1,2,3],[4,5,6],[7,8,9]]) and an index-array

arr_idx = np.array([[0,2],[1,2],[2,1]]) where each row in arr_idx corresponds to the index of arr which I want to have i.e the result should be [[1,3],[5,6],[9,8]].

Can this be done using numpy-operations? I can do it using e.g listcomprehension, but it’s some fairly large data that I have, thus if we can vectorize it, that would be better.

I have tried


result = arr[arr_idx]

which results in

array([[[1, 2, 3],
        [7, 8, 9]],
       [[4, 5, 6],
        [7, 8, 9]],
       [[7, 8, 9],
        [4, 5, 6]]])

which should’ve been array([[1,3],[5,6],[9,8]])

Asked By: CutePoison

||

Source

Answer 1

You need to use:

result = arr[np.arange(arr.shape[0])[:,None], arr_idx]

Or:

result = np.take_along_axis(arr, arr_idx, 1)

Output:

array([[1, 3],
       [5, 6],
       [9, 8]])

Answered By: mozway

Get columns of 2D ndarray defined by index' in another 2D ndarray

Question:

Answers: