编写cython函数:

Write a cython function:

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的结果:

The result of %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.take() is the fastest if both of the arguments are ndarray object. The cython version is 5x faster than list comprehension.