CODEWITHSUNDEEP
pythonnumpy
Matteo B.
Matteo B.

Reputation: 4074

Select multiple elements per row in numpy

With

ad = np.array([ 0.5,  0.8,  0.9,  0.1])
cp = np.array([[2,3,1,1,2,0],[1,0,1,3,1,2],[1,1,1,1,1,1],[0,1,2,3,2,2]])

How can I get Numpy to give me the elements of ad with the indexes of cp[0,:] as first row (the indexes are [2,3,1,1,2,0] so the first row should be [0.9, 0.1, 0.8, 0.8, 0.9, 0.5]), the elements with the indexes of cp[1,:] as second row etc.?

So the result should be:

[[0.9, 0.1, 0.8, 0.8, 0.9, 0.5],
 [0.8, 0.5, 0.8, 0.1, 0.8, 0.9],
 [0.8, 0.8, 0.8, 0.8, 0.8, 0.8],
 [0.5, 0.8, 0.9, 0.1, 0.9, 0.9]]

Preferably, of course, in an efficient and elegant way.

Upvotes: 0

Views: 211

Answers (1)

Mike Müller
Mike Müller

Reputation: 85452

NumPy arrays support broadcasting. It will broadcast ad to the needed shape. So this

>>> ad[cp]
array([[ 0.9,  0.1,  0.8,  0.8,  0.9,  0.5],
       [ 0.8,  0.5,  0.8,  0.1,  0.8,  0.9],
       [ 0.8,  0.8,  0.8,  0.8,  0.8,  0.8],
       [ 0.5,  0.8,  0.9,  0.1,  0.9,  0.9]])

will work.

Alternatively, you can use np.take():

>>> np.take(ad, cp)
array([[ 0.9,  0.1,  0.8,  0.8,  0.9,  0.5],
       [ 0.8,  0.5,  0.8,  0.1,  0.8,  0.9],
       [ 0.8,  0.8,  0.8,  0.8,  0.8,  0.8],
       [ 0.5,  0.8,  0.9,  0.1,  0.9,  0.9]])

Upvotes: 1

Related Questions