在Python中向量化平方欧几里德距离的掩码

Question

我正在运行代码以生成B中位置的掩码，该掩码比D中的位置更接近A中的位置。

N = [[0 for j in range(length_B)] for i in range(length_A)]    
dSquared = D*D

for i in range(length_A):
    for j in range(length_B):
        if ((A[j][0]-B[i][0])**2 + (A[j][1]-B[i][1])**2) <= dSquared:
            N[i][j] = 1

对于长度为数万个位置的A和B的列表，此代码需要一些时间。我很确定有办法矢量化它，但是可以使其运行得更快。谢谢。

Answer 1

您可以使用scipy\'s cdist它来进行此类距离计算，据说非常有效，如下所示 -

\n\n

from scipy.spatial.distance import cdist\n\nN = (cdist(A,B,\'sqeuclidean\') <= dSquared).astype(int)\n

\n\n\n\n

正如 中所建议的@hpaulj\'s solution，可以使用也可以使用broadcasting。现在，从问题中发布的代码来看，我们似乎正在处理Nx2成形数组。因此，我们基本上可以对第一列和第二列进行切片，并分别对它们执行广播减法。好处是我们不会去3D，因此保持内存效率，这也可能转化为性能提升。因此，平方欧氏距离的计算方式如下 -

\n\n

sq_eucl_dist = (A[:,None,0] - B[:,0])**2 + (A[:,None,1] - B[:,1])**2\n

\n\n

让我们对所有这三种方法进行计时，以进行平方欧氏距离计算。

\n\n

运行时测试 -

\n\n

In [75]: # Input arrays\n    ...: A = np.random.rand(200,2)\n    ...: B = np.random.rand(200,2)\n    ...: \n\nIn [76]: %timeit ((A[:,None,:] - B[None,:,:])**2).sum(axis=-1) # @hpaulj\'s solution\n1000 loops, best of 3: 1.9 ms per loop\n\nIn [77]: %timeit (A[:,None,0] - B[:,0])**2 + (A[:,None,1] - B[:,1])**2\n1000 loops, best of 3: 401 \xc2\xb5s per loop\n\nIn [78]: %timeit cdist(A,B,\'sqeuclidean\')\n1000 loops, best of 3: 249 \xc2\xb5s per loop\n

\n