CODEWITHSUNDEEP
pythonnumpy
FenryrMKIII
FenryrMKIII

Reputation: 1198

Numpy array row indexing through an array of indices

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 ?

Upvotes: 1

Views: 121

Answers (1)

akuiper
akuiper

Reputation: 214927

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

Upvotes: 1

Related Questions