NumPy：利用1D numpy数组来切片列表

Question

如何提取列表中与 1D numpy.ndarray 中包含的索引对应的元素？

以下是示例：

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

list_data[arr_index]  # FAILS

我希望能够检索与arr_index相对应的list_data元素。

- tchakravarty

Answer 1

你可以使用 numpy.take 函数。

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

样例运行 -

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])

- Divakar

---

@Kasra，你所说的“简单索引”是什么意思？ - tchakravarty

@Kasra 显然，简单的索引期望第一个输入为NumPy数组。 - Divakar

@fgnu：它被转换为ndarray。您可以在source中看到，对于列表输入，“take”最终调用“_wrapit”，后者调用“asarray”。就我个人而言，我要么使用列表推导式，要么在最早的可能点将列表转换为数组。 - user2357112

@fgnu：如果你只使用它一次，列表推导式不必遍历列表的每个元素。哪种选项最快取决于你希望获取多少项。你也可以尝试类似map(list_data.__getitem__, arr_index)这样的方法，看看是否有帮助。 - user2357112

@user2357112，我预计要取长列表（20K个元素）的约3％的元素。您建议使用哪种习语更具性能？ - tchakravarty

Answer 2

我刚做了一些计时测试：

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

列表推导式明显是最快的。虽然使用数组索引速度更快，但创建数组的开销很大。

转换为对象 dtype 数组更快。我有点惊讶，但这一定是因为它可以在不解析的情况下进行转换：

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

- hpaulj

Answer 3

或者

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]

注意

可以用 list_data = range(1, 100) 替换 list_data = list(range(1, 100) 可以用 arr_index = np.array([18, 55, 22]) 替换 arr_index = np.asarray([18, 55, 22])

- farhawa

---

在py3中，除非您立即进行迭代，否则需要使用list(range())。 - hpaulj