CODEWITHSUNDEEP
pythonarraysnumpy
Nico G.
Nico G.

Reputation: 517

2D to 1D numpy array with indices of column for each row

I'm trying to index a column for each row in a matrix.

Suppose I have a numpy array A with the shape (n,m).
I also have a numpy array B with the shape (n,) containing integers between 0 and m, so they can be used as indices for the 2nd axis of A.

I want to get a numpy array C with shape (n,) with C[i] = A[i,B[i]], so every row of A yields one value based on the index in B.

Surely I could use this last expression in a for loop or list comprehension, but how would I do it using the efficiency of numpy?

My first intuition was C = A[:,B] but this clearly yields something else. (shape (n,n))

Upvotes: 0

Views: 977

Answers (1)

myz540
myz540

Reputation: 559

You are close, try:

C = A[range(len(B)), B]

This should yield C[i] = A[i,B[i]]

Upvotes: 1

Related Questions