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]

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

Answer 2

您的方法是二次的，这里是排序数组的单程线性时间算法。您只需在正确的时间运行两个阵列。

def prox(a,b,dx):
    ia=ib=ir=0
    res=zeros(a.size,int32)
    while ia<a.size and ib<b.size:
        if abs(a[ia]-b[ib])<dx: 
            res[ir]=ia
            ir += 1
            ia += 1
        elif a[ia]>b[ib] :
               ib += 1
        else : ia += 1
    return res[:ir]

您可以使用Numba编译此代码以进一步提高性能。

测试：

a=rand(1000)
b=rand(1000)
a.sort()
b.sort()

In [10]:   prox(a,b,1e-5)
Out[10]: 
array([ 35,  90, 159, 165, 174, 252, 276, 380, 383, 467, 508, 515, 641,
       658, 705, 711, 728, 814, 857, 871, 907, 945])

In [11]: %timeit prox(a,b,1e-4)
100 loops, best of 3: 6.23 ms per loop
In [12]: prox2=numba.jit(prox)
In [13]: %timeit prox2(a,b,1e-4)
10000 loops, best of 3: 19.1 µs per loop

Answer 3

这并没有利用数据的排序特性，因此它没有线性时间复杂度（虽然我怀疑运行时确实受益于它的排序，缓存方式），但nlogn并不坏，而且肯定很难在简单性和良好测试方面击败：

from scipy.spatial import cKDTree
print(cKDTree(a[:, None]).query_ball_point(b[:, None], dx))

Answer 4

当a和b是排序数组时，此排序的性质可能会被{em>滥用并带有np.searchsorted。基本思想是我们得到b的索引，其中可以放置a中的每个元素，以便维护排序的顺序。通过这种方式，我们可以知道b中特定bin的右侧边界，其中可以放置a中的每个元素。要获取同一个bin的左侧边界，只需从先前找到的那些索引中减去1。然后，获取每个a与左边界和右边界之间的差异，看看是否在阈值范围内，如果是，请将其视为有效索引。

对于极端情况会有一些额外的工作，即当a中的元素大于最高b时，以及当元素小于最低b时。如果我们使用np.searchsorted的{{1}}搜索选项，我们只需将其剪切到最小'left'即可找到正确的边界，以便可以在整个范围内使用这些相同的索引整个阵列一气呵成。因此，实现看起来像这样 -

运行时测试 -

def find_all_close_searchsorted(a, b):
    lidx = np.searchsorted(b,a,'left').clip(min=1,max=b.size-1)
    close_mask = (np.abs(b[lidx] - a) <= dx) | (np.abs(b[lidx-1] - a) <= dx)
    return np.nonzero(close_mask)[0]

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

4 个答案: