Question

我有一个名为list的常规a和一个NumPy索引b数组。
（不，我不可能将a转换为NumPy数组。）

对我有什么办法和＃34; a[b]＆＃34;有效率的？需要明确的是，这意味着由于其性能影响，我不想提取int中的每个b。

（是的，这是我的代码中的瓶颈。这就是我开始使用NumPy数组的原因。）

Answer 1

a = list(range(1000000))
b = np.random.randint(0, len(a), 10000)

%timeit np.array(a)[b]
10 loops, best of 3: 84.8 ms per loop

%timeit [a[x] for x in b]
100 loops, best of 3: 2.93 ms per loop

%timeit operator.itemgetter(*b)(a)
1000 loops, best of 3: 1.86 ms per loop

%timeit np.take(a, b)
10 loops, best of 3: 91.3 ms per loop

我对numpy.take()寄予厚望，但远非最佳。我也尝试了一些Numba解决方案，他们产生了类似的时间 - 大约92毫秒。

因此，一个简单的列表理解与这里的最佳理解并不相同，但operator.itemgetter()获胜，至少对于这些数量级的输入大小而言。

Answer 2

编写一个cython函数：

import cython
from cpython cimport PyList_New, PyList_SET_ITEM, Py_INCREF

@cython.wraparound(False)
@cython.boundscheck(False)
def take(list alist, Py_ssize_t[:] arr):
    cdef:
        Py_ssize_t i, idx, n = arr.shape[0]
        list res = PyList_New(n)
        object obj

    for i in range(n):
        idx = arr[i]
        obj = alist[idx]
        PyList_SET_ITEM(res, i, alist[idx])
        Py_INCREF(obj)

    return res

％timeit的结果：

import numpy as np

al= list(range(10000))
aa = np.array(al)

ba = np.random.randint(0, len(a), 10000)
bl = ba.tolist()

%timeit [al[i] for i in bl]
%timeit np.take(aa, ba)
%timeit take(al, ba)

1000 loops, best of 3: 1.68 ms per loop
10000 loops, best of 3: 51.4 µs per loop
1000 loops, best of 3: 254 µs per loop

如果两个参数都是ndarray对象，则

numpy.take()是最快的。 cython版本比列表理解快5倍。

使用NumPy索引数组切片Python列表 - 任何快速方式？

2 个答案: