比较两个不同长度的 numpy 数组

Question

我需要找到另一个数组中一个数组的第一个小于或等于元素出现的索引。 一种有效的方法是：

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”循环很慢，我的数组可能有数百万个元素。 这有什么麻木的技巧吗？ 这不起作用，因为它们的数组长度不同：

indices = numpy.where(a<=b) #XXX: raises an exception

谢谢！

Answer 1

这可能是一种特殊情况，但您应该能够使用 numpy digitize 。 这里的警告是 bin 必须单调减少或增加。

>>> 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

看起来加速了两个数量级，但这在很大程度上取决于 bin 的数量。 你的时间会有所不同。

---

为了与杰米的回答进行比较，我对以下两段代码进行了计时。 因为我主要想关注搜索searchsorted与digitize的速度， searchsorted我searchsorted缩减了 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的简单语法几乎同样好。

Answer 2

这很混乱，但它有效：

>>> idx = np.argsort(a)
>>> np.take(idx, np.searchsorted(a, b, side='right', sorter=idx)-1)
array([0, 1, 1, 1, 2, 2, 2, 2, 2, 3], dtype=int64)

如果你的数组总是排序的，你应该能够摆脱argsort调用。