Efficient implementation (vectorized) in numpy

Question

I have a matrix in numpy

A = [[1,2,3],
     [4,5,6],
     [7,8,9]]

and a plain array

B = [2,1,0]

I want to compress the matrix into a single column such that for the ith row I want to pick the B[i]th column and put it in the column vector

So I want an elegant/vectorized implementation to get

Ans = [3,5,7]

score 1 · Accepted Answer · answered Dec 01 '17 at 15:25

1

Try this:

In [32]: A[np.arange(A.shape[0]), B]
Out[32]: array([3, 5, 7])

answered Dec 01 '17 at 15:25

I tried to use index slicing instead of arange, something like A[:,B]/A[:,B[:]] ofcourse didnt work! This is elegant thanks!! – Naman Dec 01 '17 at 16:00
@Naman, glad it helps :) – MaxU - stop genocide of UA Dec 01 '17 at 16:02

1 Answers1