Selecting multiple elements from an 2d array using a 1d array in numpy

Question

I have two arrays in numpy. The first is a 2d array, which can be thought of as a list of vectors. The second is a 1d array, which can be thought of as a list of indices into the 2d array.

I want to select elements of the 2d array using the indices of the 1d array. Right now I have been doing

        z=rnd.rand(2,10) # a list of 2d vectors of length 10
        z_idx=rnd.randint(2,size=z.shape[1]) #indices selecting a dimension of the 2d vector

        result=np.array([z[z_idx[i],i] for i in xrange(len(z_idx))])

But this is very slow.

Is there a better way to do this in numpy?

Answer 1

Probably the simplest method:

result = z[z_idx].diagonal()

Maybe a little more efficient would be to use arange :

result = z[z_idx, np.arange(z_idx.size)]

More appropriate but equivalent is np.indices :

result = z[z_idx, np.indices(z_idx.shape)[0]]