Index column values in 2D array using array of ind

Question

This question already has answers here:

using an numpy array as indices of the 2nd dim of another array? [duplicate] (2 answers)

Closed 9 days ago.

I have the following array:

import numpy as np 

print(A)
array([[ 0,  1,  4,  5,  8,  7],
       [ 5,  3,  4,  1,  8, 11],
       [ 2,  7,  5,  3,  4,  1],
       [ 2,  8,  8,  1, 10,  1],
       [ 2, 14,  8,  6,  5,  3]])

And I need to the values A corresponding to these column indices:

b = np.array([5, 0, 3, 4, 4])

Expected output:

array([ 7,  5,  3, 10,  5])

Thanks in advance.

Answer 1

You can use advanced indexing. You need to define an indexing array across the first axis, so that both indexing arrays are broadcast together and each column index refers to a specific row. In this case you just want an np.arange to index on the rows:

A[np.arange(A.shape[0]), b]
# array([ 7,  5,  3, 10,  5])