python - 基于类别更有效地分配 numpy 数组元素

Question

我有一个一维数组，其值根据自然中断算法分组(Jenks，感兴趣的人)。我想根据初始数组的组创建另一个包含元素的数组。这是一个例子:

import numpy as np
arr = np.array([1, 42, 1, 1, 2, 43, 2, 3, 44, 41, 42])
bins = np.array([1.0, 3.0, 44.0])
rank = 1 * (arr == bins[0])
# Next two lines to be optimized
for i in xrange(len(bins) - 1):
    rank[(arr > bins[i]) & (arr <= bins[i+1])] = bins[i+1]
print rank
# [ 1 44  1  1  3 44  3  3 44 44 44]

解释一下，这会经历bins中的间隔(a, b]，并将rank中所有元素的b值分配给arr间隔。这个代码片段产生了我想要的输出，但是当 bins 很大时，循环可能需要很长时间。是否有更有效的方法来做到这一点，可能使用一些 numpy 工具？

Answer 1

我认为您正在寻找 numpy.searchsorted 函数。我认为它会比您自己编写的几乎任何循环都快。

arr = np.array([1, 42, 1, 1, 2, 43, 2, 3, 44, 41, 42])
bins = np.array([1.0, 3.0, 44.0])
rank = bins[np.searchsorted(bins, arr)]
print rank
[  1.  44.   1.   1.   3.  44.   3.   3.  44.  44.  44.]

还有 numpy.histogram 您可能也想看看。它更容易理解，但在内部它也使用 searchsorted。

更新:我将示例数组的计时与 timeit 进行了比较。

%timeit rank = bins[np.searchsorted(bins, arr)]
1000000 loops, best of 3: 1.51 µs per loop

这与您的原始循环相比效果很好:

%timeit for i in xrange(len(bins) - 1): rank[(arr > bins[i]) & (arr <= bins[i+1])] = bins[i+1]
100000 loops, best of 3: 8.78 µs per loop