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

【问题标题】：Comparing two numpy arrays of different length比较两个不同长度的numpy数组
【发布时间】：2020-05-21 05:57:07
【问题描述】：

我需要在另一个数组中找到一个数组的第一个小于或等于元素的索引。一种可行的方法是：

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]

indices 的值是 [0, 1, 1, 1, 2, 2, 2, 2, 2, 3]，这正是我所需要的。当然，问题是python“for”循环很慢，我的数组可能有数百万个元素。这有什么麻木的技巧吗？这不起作用，因为它们的数组长度不同：

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

谢谢！

【问题讨论】：

你能为a和b提供合理的尺寸吗？方法的时机受此影响很大。

标签： python numpy

【解决方案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

看起来速度提高了两个数量级，但这在很大程度上取决于垃圾箱的数量。您的时间安排会有所不同。

为了与 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 的原始速度对于大型情况几乎总是更快，但digitize 的简单语法几乎在 bin 数量较少的情况下也一样好。

【讨论】：

+1 - 我认为这对于问题指定的所有情况都是正确的：“我需要一个数组的元素的第一个小于或等于的索引另一个数组。”
我不确定np.digitize，但np.searchsorted 进行二分查找，所以它的强度应该是大a 数组，而不是大b 数组。设置a = np.arange(1000)[::-1] 和b = np.random.randint(0, 1000, 1E6) 和时间似乎证实了这一点。
@Jamie 这也是我发现的，所以我暂时删除了你的例子的时间——在没有更多例子的情况下比较这两者是不公平的。
@Jamie 用一些相关的时间更新了这个问题，如果你发现任何明显的缺陷，请告诉我。我留下了np.take，因为它的缩放可能是N，而searchsorted 将是主要的N log(N) 操作。
我希望我能再给你一个 +1 的精彩时间分析！

【解决方案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 调用。

【讨论】：