Numpy: How to select row entries in a 2d array by column vector

Question

How can I retrieve a column vector from a 2d array given an indicator column vector?

Suppose I have

X = np.array([[1, 4, 6],
              [8, 2, 9],
              [0, 3, 7],
              [6, 5, 1]])

and

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

Is there an elegant way to get from X and S the result array([1, 9, 3, 1]), which is equivalent to

np.array([x[s] for x, s in zip(X, S)])

Answer 1

You can achieve this using np.take_along_axis:

>>> np.take_along_axis(X, S[..., None], axis=1)
array([[1],
       [9],
       [3],
       [1]])

You need to make sure both array arguments are of the same shape (or broadcasting can be applied), hence the S[..., None] broadcasting.

Of course your can reshape the returned value with a [:, 0] slice.

>>> np.take_along_axis(X, S[..., None], axis=1)[:, 0]
array([1, 9, 3, 1])

---

Alternatively you can just use indexing with an arangement:

>>> X[np.arange(len(S)), S[np.arange(len(S))]]
array([1, 9, 3, 1])

I believe this is also equivalent to np.diag(X[:, S]) but with unnecessary copying...

Answer 2

For 2d arrays

# Mention row numbers as one list and S which is column number as other
X[[0, 1, 2, 3], S]

# more general
X[np.indices(S.shape), S]

indexing_basics