CODEWITHSUNDEEP
pythonarraysnumpy
rapto
rapto

Reputation: 419

is it possible to use numpy.take to pick a list of indices?

I want to take some indices from the elements of a numpy.array. This is what I've got so far:

import numpy as np

a = np.array(range(20))
a.shape=4,5
#~ [[ 0  1  2  3  4]
#~ [ 5  6  7  8  9]
#~ [10 11 12 13 14]
#~ [15 16 17 18 19]]

idxs = np.array([(4,0),(3,1),(2,1),(0,3)])
ret = a.take(idxs,1).diagonal().transpose()
print ret

#~ [[  4.   0.]
#~ [  8.   6.]
#~ [ 12.  11.]
#~ [ 15.  18.]]

Is is possible to get this result in a simpler (or faster) way?

Upvotes: 0

Views: 273

Answers (1)

HYRY
HYRY

Reputation: 97331

you can create an array for axis 0 index:

a[np.arange(4)[:,None], idxs]

Upvotes: 3

Related Questions