更快速地根据阈值拆分numpy数组

Question

假设我有一个随机的numpy数组:

X = np.arange(1000)

和一个门槛:

thresh = 50

我想要分割X在两个分区X_l和X_r以这样的方式,在每一个元件X_l是小于或等于thresh而在X_r每个元素是大于thresh.之后,这两个分区被赋予递归函数.

使用numpy我创建一个布尔数组,我用它来分区X:

Z = X <= thresh
X_l, X_r = X[Z == 0], X[Z == 1]
recursive_call(X_l, X_r)

这样做了几次,有没有办法让事情变得更快？是否可以避免在每次调用时创建分区的副本？

Answer 1

X[~Z]比X[Z==0]以下更快:

In [13]: import numpy as np

In [14]: X = np.random.random_integers(0, 1000, size=1000)

In [15]: thresh = 50

In [18]: Z = X <= thresh

In [19]: %timeit X_l, X_r = X[Z == 0], X[Z == 1]
10000 loops, best of 3: 23.9 us per loop

In [20]: %timeit X_l, X_r = X[~Z], X[Z]
100000 loops, best of 3: 16.4 us per loop

您是否已经分析过确定这确实是代码中的瓶颈？如果您的代码仅花费1%的时间进行此拆分操作,那么优化此操作的程度将对整体性能产生不超过1%的影响.

通过重新思考算法或数据结构而不是优化这一操作,您可能会受益更多.如果这确实是瓶颈,你可以通过在C或Cython中重写这段代码来做得更好 ......

当你有大小1000 numpy的阵列,有一个机会,使用Python列表/套/类型的字典可能会更快.在阵列非常大之前,NumPy阵列的速度优势有时并不明显.您可能希望用纯Python重写代码,并使用timeit对这两个版本进行基准测试.

嗯,让我重新说一下.它不是真正的数组大小使NumPy更快或更慢.它只是拥有小型NumPy数组有时表明你正在创建许多小的NumPy数组,并且创建NumPy数组要比创建Python列表慢得多:

In [21]: %timeit np.array([])
100000 loops, best of 3: 4.31 us per loop

In [22]: %timeit []
10000000 loops, best of 3: 29.5 ns per loop

In [23]: 4310/295.
Out[23]: 14.610169491525424

此外,当您使用纯Python进行编码时,您可能更有可能使用没有直接NumPy等效项的dicts和集合.这可能会导致您使用更快的替代算法.