我需要找到另一个数组中一个数组中第一个小于或等于元素的索引。一种有效的方式是这样的:
import numpy
a = numpy.array([10,7,2,0])
b = numpy.array([10,9,8,7,6,5,4,3,2,1])
indices = [numpy.where(a<=x)[0][0] for x in b]
索引的值为[0、1、1、1、2、2、2、2、2、3],这是我需要的。当然,问题是python“for”循环很慢,我的数组可能有数百万个元素。是否有任何numpy技巧?这是行不通的,因为它们的数组长度不同:
indices = numpy.where(a<=b) #XXX: raises an exception
谢谢!
最佳答案
这可能是一种特殊情况,但是您应该可以使用numpy digitize。需要注意的是,垃圾箱必须单调减少或增加。
>>> import numpy
>>> a = numpy.array([10,7,2,0])
>>> b = numpy.array([10,9,8,7,6,5,4,3,2,1])
>>> indices = [numpy.where(a<=x)[0][0] for x in b]
[0, 1, 1, 1, 2, 2, 2, 2, 2, 3]
>>> numpy.digitize(b,a)
array([0, 1, 1, 1, 2, 2, 2, 2, 2, 3])
定时测试的设置:
a = np.arange(50)[::-1]
b = np.random.randint(0,50,1E3)
np.allclose([np.where(a<=x)[0][0] for x in b],np.digitize(b,a))
Out[55]: True
一些时间安排:
%timeit [np.where(a<=x)[0][0] for x in b]
100 loops, best of 3: 4.97 ms per loop
%timeit np.digitize(b,a)
10000 loops, best of 3: 48.1 µs per loop
看起来速度加快了两个数量级,但是,这在很大程度上取决于存储箱的数量。您的时间安排会有所不同。
为了与Jamie的答案进行比较,我对以下两个代码进行了计时。因为我主要想关注
searchsorted
和digitize
的速度,所以我简化了Jamie的代码。相关的块在这里:a = np.arange(size_a)[::-1]
b = np.random.randint(0, size_a, size_b)
ja = np.take(a, np.searchsorted(a, b, side='right', sorter=a)-1)
#Compare to digitize
if ~np.allclose(ja,np.digitize(b,a)):
print 'Comparison failed'
timing_digitize[num_a,num_b] = timeit.timeit('np.digitize(b,a)',
'import numpy as np; from __main__ import a, b',
number=3)
timing_searchsorted[num_a,num_b] = timeit.timeit('np.take(a, np.searchsorted(a, b, side="right", sorter=a)-1)',
'import numpy as np; from __main__ import a, b',
number=3)
这有点超出了我有限的matplotlib功能,因此可以在DataGraph中完成。我已经绘制了
timing_digitize/timing_searchsorted
的对数比率,因此大于零searchsorted
的值更快,而小于零digitize
的值更快。颜色还给出了相对速度。例如,显示在右上角(a = 1E6,b = 1E6)digitize
比searchsorted
慢约300倍,而对于较小的digitize
可以快10倍。黑线大致是收支平衡点:看起来原始速度
searchsorted
在大多数情况下几乎总是更快,但是如果bin的数量很少,那么digitize
的简单语法几乎一样好。