在Python中矢量化平方欧几里得距离的掩码答案

【问题标题】：Vectorize mask of squared euclidean distance in Python在Python中矢量化平方欧几里得距离的掩码
【发布时间】：2016-08-06 07:59:40
【问题描述】：

我正在运行代码来生成 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 列表，此代码需要一段时间。我很确定有一种方法可以对其进行矢量化，以使其运行得更快。谢谢。

【问题讨论】：

不应该是N[j][i] = 1吗？
@Divakar – 你是对的，我更改了要在此处发布的项目名称以使其更简单，并且感到困惑。我已相应地更正了代码。

标签： python numpy scipy vectorization euclidean-distance

【解决方案1】：

只要您的矩阵 N 可能是稀疏的，scipy.spatial.cKDTree 将提供比任何基于计算所有距离蛮力的方法更好的时间复杂度：

cKDTree(A).sparse_distance_matrix(cKDTree(B), max_distance=D)

【讨论】：

【解决方案2】：

您可以使用scipy's cdist，这对于此类距离计算来说应该非常有效，就像这样 -

from scipy.spatial.distance import cdist

N = (cdist(A,B,'sqeuclidean') <= dSquared).astype(int)

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

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

让我们为所有这三种方法计算平方欧几里得距离的时间。

运行时测试-

In [75]: # Input arrays
    ...: A = np.random.rand(200,2)
    ...: B = np.random.rand(200,2)
    ...: 

In [76]: %timeit ((A[:,None,:] - B[None,:,:])**2).sum(axis=-1) # @hpaulj's solution
1000 loops, best of 3: 1.9 ms per loop

In [77]: %timeit (A[:,None,0] - B[:,0])**2 + (A[:,None,1] - B[:,1])**2
1000 loops, best of 3: 401 µs per loop

In [78]: %timeit cdist(A,B,'sqeuclidean')
1000 loops, best of 3: 249 µs per loop

【讨论】：

cdist.asType(bool) 与输出中的 np.nonzero() 相结合，由于输出非常稀疏，导致性能大幅提升和内存减少。感谢您的帮助，非常感谢。

【解决方案3】：

我赞同上面使用 Numpy 的建议。循环代码对 A 的索引也比它需要的要多得多。你可以使用类似的东西：

import numpy as np

dimension = 10000
A = np.random.rand(dimension, 2) + 0.0
B = np.random.rand(dimension, 2) + 1.0
N = []
d = 1.0

for i in range(len(A)):
    distances = np.linalg.norm(B - A[i,:], axis=1)
    for j in range(len(distances)):
        if distances[j] <= d:
            N.append((i,j))

print(len(N))

要从纯 Python 中获得不错的性能是相当困难的。我还要指出，更多维数组解决方案将需要...很多...内存。

【讨论】：

【解决方案4】：

使用二维数组索引更容易可视化这段代码：

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

dist 表达式看起来像

((A[j,:] - B[i,:])**2).sum(axis=1)

对于 2 个元素，这个数组表达式可能不会更快，但它应该可以帮助我们重新思考问题。

我们可以解决i,j、outter的广播问题

A[:,None,:] - B[None,:,:]  # 3d difference array

dist=((A[:,None,:] - B[None,:,:])**2).sum(axis=-1)  # (lengthA,lengthB) array

将此与dSquared 进行比较，并使用生成的布尔数组作为掩码将N 的元素设置为1：

N = np.zeros((lengthA,lengthB))
N[dist <= dSquared] = 1

我没有测试过这段代码，所以可能有错误，但我认为基本的想法就在那里。并且可能足以让您了解其他情况的详细信息。

【讨论】：