Numpy array row indexing through an array of indices

Question

Imagine I have the following (3,2) Numpy array A

A = np.array([[1, 2],
 [3, 4],
 [5, 6]])

and I would like to index this array A column-wise thanks to the indices available in numpy Array B :

B = np.array([[1,0],[2,0]])

So I would like to take rows 1 and 0 in column 0 and rows 2 and 0 in column 1 to obtain :

C = np.array([[3,1],[6,2]])

What is an efficient way to do that ?

Answer 1

You can construct the column index with np.arange(A.shape[1]), transpose B so it broadcasts with the column index correctly, and then extract the elements with advanced indexing:

A = np.array([[1, 2],
 [3, 4],
 [5, 6]])
B = np.array([[1,0],[2,0]])
A[B.T, np.arange(A.shape[1])].T
#array([[3, 1],
# [6, 2]])

Row index:

B.T
# V second column row index
#array([[1, 2],
# [0, 0]])
# ^ first column row index

Column index:

np.arange(A.shape[1])
# array([0, 1])