CODEWITHSUNDEEP
pythonarraysnumpy
Leon palafox
Leon palafox

Reputation: 2785

Index a numpy array using a list in a pythonic way

Imagine that I have a python array like

array = [[2,3,4],[5,6,7],[8,9,10]]

And a list

list = [0,2,1]

I basically want a one liner to extract the indexed elements from the array given by the list

For example, with the given array and list:

result = [2,7,9]

My kneejerk option was

result = array[:, list]

But that did not work

I know that a for cycle should do it, I just want to know if there is some indexing that might do the trick

Upvotes: 2

Views: 74

Answers (2)

Ashwini Chaudhary
Ashwini Chaudhary

Reputation: 250941

Something like this?

In [24]: a
Out[24]: 
array([[ 2,  3,  4],                                                                              
       [ 5,  6,  7],                                                                              
       [ 8,  9, 10]])                                                                             

In [25]: lis
Out[25]: [0, 2, 1]

In [26]: a[np.arange(len(a)), lis]                                                              
Out[26]: array([2, 7, 9])

Upvotes: 2

Kiwi
Kiwi

Reputation: 2816

Use enumerate to create row indices and unzip (zip(*...)) this collection to get the row indices (the range [0, len(list))) and the column indices (lis) :

a[zip(*enumerate(lis))]

Upvotes: 0

Related Questions