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

时间:2015-03-12 18:16:45

标签: python arrays numpy

我有一个单维数组,其值根据自然中断算法(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工具?

2 个答案:

答案 0 :(得分:1)

如果没有加速,循环可以简化:

In [132]: rank1=np.ones(arr.shape,dtype=bins.dtype)

In [133]: for i in range(len(bins)-1):
    rank1[arr>bins[i]]=bins[i+1]
   .....:     

In [134]: rank1
Out[134]: array([  1.,  44.,   1.,   1.,   3.,  44.,   3.,   3.,  44.,  44.,  44.])

但这是使用数组操作生成相同内容的方法。

In [135]: rank2=bins[np.sum(arr[None,:]>bins[:,None],axis=0)]

In [136]: rank2
Out[136]: array([  1.,  44.,   1.,   1.,   3.,  44.,   3.,   3.,  44.,  44.,  44.])

更好吗?它可能更加模糊。在你的小例子的时间测试中,这两个版本需要相同的时间,大约是你版本速度的2倍。

答案 1 :(得分: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