NumPy: 1D numpy array to slice a list

Question

How can I extract the elements of a list corresponding to the indices contained in a 1D numpy.ndarray ?

Here is an example:

list_data = list(range(1, 100))
arr_index = np.asarray([18, 55, 22])
arr_index.shape

list_data[arr_index]  # FAILS

I want to be able to retrieve the elements of list_data corresponding to arr_index .

Answer 1

You can use numpy.take -

import numpy as np
np.take(list_data,arr_index)

Sample run -

In [12]: list_data = list(range(1, 20))

In [13]: list_data
Out[13]: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19]

In [14]: arr_index = np.asarray([3, 5, 12])

In [15]: np.take(list_data,arr_index)
Out[15]: array([ 4,  6, 13])

Answer 2

OR

import numpy as np
list_data = list(range(1, 100))
arr_index = np.asarray([18, 55, 22])
arr_index.shape

new_ = [list_data[i] for i in arr_index]

>> [19, 56, 23]

Note

list_data = list(range(1, 100) can be replaced by list_data = range(1, 100)

arr_index = np.asarray([18, 55, 22]) can be replaced by arr_index = np.array([18, 55, 22])

Answer 3

I just did some timing tests:

In [226]: ll=list(range(20000))    
In [227]: ind=np.random.randint(0,20000,200)

In [228]: timeit np.array(ll)[ind]
100 loops, best of 3: 3.29 ms per loop

In [229]: timeit np.take(ll,ind)
100 loops, best of 3: 3.34 ms per loop

In [230]: timeit [ll[x] for x in ind]
10000 loops, best of 3: 65.1 µs per loop

In [231]: arr=np.array(ll)
In [232]: timeit arr[ind]
100000 loops, best of 3: 6 µs per loop

The list comprehension clearly is the winner. Indexing an array is clearly faster, but the overhead of creating that array is substantial.

Converting to an object dtype array is faster. I'm a little surprised, but it must be because it can convert without parsing:

In [236]: timeit np.array(ll,dtype=object)[ind].tolist()
1000 loops, best of 3: 1.04 ms per loop