Rakshit Kothari
Rakshit Kothari

Reputation: 416

How to access entries of a list using a numpy array without using a for loop

Given a list L = [0,1,2,3,4,5,6,7,8,9]. What's the best way to access/extract elements where their indices are given by a numpy array? Let nArr=np.array([0,-1,2,6]).

The resulting output should be another list P = [0,9,2,6].

It is clear that when the elements are uniform in shape, we can simply cast it into another numpy array, but what if it isn't? For example, M = [np.random.rand(5,5), np.random.rand(1)].

Upvotes: 1

Views: 885

Answers (2)

user6655984
user6655984

Reputation:

To summarize the comments: lists do not support indexing by an array, like L[nArr] where nArr is an array of indexes. One normally uses list comprehension, [L[i] for i in nArr]. But if you want to, you can cast the list to a NumPy array of objects, which can then be indexed and sliced as any other NumPy array:

np.array(L, dtype=object)[nArr]

If you want a list returned, you can do:

np.array(L, dtype=object)[nArr].tolist()

But that's not nearly as pythonic as list comprehension, requires more memory, and very likely doesn't improve the speed.

Upvotes: 1

hpaulj
hpaulj

Reputation: 231335

Stock Python has a convenience class, itemgetter:

In [27]: from operator import itemgetter
In [28]: L = [0,1,2,3,4,5,6,7,8,9]
In [29]: nArr=np.array([0,-1,2,6])
In [30]: itemgetter(*nArr)
Out[30]: operator.itemgetter(0, -1, 2, 6)
In [31]: itemgetter(*nArr)(L)
Out[31]: (0, 9, 2, 6)

Internally it does something equivalent to the list comprehension:

In [33]: [L[x] for x in nArr]
Out[33]: [0, 9, 2, 6]

So it isn't a fast operation like the array indexing (look at its code). It may be most useful as a way of doing sort or other operations where you'd like to define a key function that fetches multiple values.

https://stackoverflow.com/a/47585659/901925

Make a random nested list:

In [54]: arr = np.random.randint(0,10,(4,4))
In [55]: L = arr.tolist()
In [56]: L
Out[56]: [[9, 5, 8, 4], [1, 5, 5, 8], [8, 0, 5, 8], [1, 4, 5, 1]]

lexical sort by 'column':

In [57]: sorted(L)
Out[57]: [[1, 4, 5, 1], [1, 5, 5, 8], [8, 0, 5, 8], [9, 5, 8, 4]]

lexical sort by 'columns' 2 and 1 (in that order):

In [59]: sorted(L, key=itemgetter(2,1))
Out[59]: [[8, 0, 5, 8], [1, 4, 5, 1], [1, 5, 5, 8], [9, 5, 8, 4]]

Upvotes: 2

Related Questions