numpy - select multiple elements from each row of an array

Question

I need to select multiple different values from each row of a 2D array.

A = np.array([[ 1, 2, 3, 4],
              [ 5, 6, 7, 8],
              [ 9,10,11,12])
A[something]

>>> np.array([[ 1, 2],
              [ 6, 7],
              [11,12]])

I know I can create a boolean array the same shape as A and set each element in a for loop, but I'm hoping come up with a better solution.

Answer 1

You can try the following:

import numpy as np
A = np.array([[ 1, 2, 3, 4],
              [ 5, 6, 7, 8],
              [ 9,10,11,12]])
i = [[0],[1],[2]]
j = [[0,1], [1,2],[2,3]]
B = A[i,j]
print(B)
#Prints
[[ 1  2]
 [ 6  7]
 [11 12]]