更快速地根据阈值拆分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和集合。 这可能会导致您使用更快的替代算法。

Answer 2

您的阵列是否已经排序？ 在您的例子中，你使用的arange ，这是排序的，所以没有必要做布尔索引，你可以简单地在适当的位置切开你的一半阵列。 这样可以避免使用“高级索引”，因此您无需复制数组。

X = np.arange(0, 2*thresh)
i = X.searchsorted(thresh, side='right') # side='right' for `<=`
X_l, X_r = X[:i], X[i:]

这为排序的数组节省了大量时间，但显然不会起作用：

thresh = 500
X = np.arange(2*thresh)

%%timeit
i = X.searchsorted(thresh, side='right')
X_l, X_r = X[:i], X[i:]
100000 loops, best of 3: 5.16 µs per loop

%%timeit
Z = X <= thresh                         
X_l, X_r = X[Z], X[~Z]
100000 loops, best of 3: 12.1 µs per loop