查找一个数组中哪些元素与另一个元素中的任何元素接近的最有效方法是什么？

Question

我有两个1维numpy.ndarray对象,并想知道第一个数组中的哪些元素在第二个数组中dx的任何元素内.

我现在拥有的是什么

# setup
numpy.random.seed(1)
a = numpy.random.random(1000)  # create one array
numpy.random.seed(2)
b = numpy.random.random(1000)  # create second array
dx = 1e-4  # close-ness parameter

# function I want to optimise
def find_all_close(a, b):
    # compare one number to all elements of b
    def _is_coincident(t):
        return (numpy.abs(b - t) <= dx).any()
    # vectorize and loop over a
    is_coincident = numpy.vectorize(_is_coincident)
    return is_coincident(a).nonzero()[0]

返回timeit结果如下

10 loops, best of 3: 16.5 msec per loop

什么是优化的最佳方式find_all_close的功能,特别是如果a与b保证是float数组按升序排序,当他们获得通过到find_all_close,可能与用Cython或类似的？

在实践中,我正在使用10,000到100,000个元素(或更大)的数组,并在几百个不同的b数组上运行整个操作.

Answer 1

最简单的方法是对于第一个数组中的每个元素，对第二个数组进行两次二分搜索，以查找第一个数组中的元素下方最多 dx 和上方最多 dx 的元素。这是线性时间：

left = np.searchsorted(b, a - dx, 'left')
right = np.searchsorted(b, a + dx, 'right')
a[left != right]

线性算法有两个指向第二个数组的指针，当您迭代第一个数组中的元素时，它们会跟踪移动窗口。