Vectorized way of accessing row specific elements

Question

I have a 2-D NumPy array and a set of indices the size of which is the first dimension of the NumPy array.

X = np.random.rand(5, 3)
a = np.random.randint(0, 3, 5)

I need to do something like

for i, ind in enumerate(a):
    print X[i][ind]

Is there a vectorized way of doing this?

Answer 1

Here you go:

X = np.random.rand(5, 3)
a = np.random.randint(0, 3, 5)

In [12]: X[np.arange(a.size), a]
Out[12]: array([ 0.99653335,  0.30275346,  0.92844957,  0.54728781,  0.43535668])
In [13]: for i, ind in enumerate(a):
            print X[i][ind]
#   ....:
#0.996533345844
#0.30275345582
#0.92844956619
#0.54728781105
#0.435356681672

I'm assuming here that you don't need each value on a separate line and just want to extract the values.