您可以使用
scipy的cdist
进行距离计算,据说这种方法非常高效,例如:
from scipy.spatial.distance import cdist
N = (cdist(A,B,'sqeuclidean') <= dSquared).astype(int)
如
@hpaulj的解决方案
所建议的那样,我们也可以使用。现在,从问题中发布的代码来看,似乎我们正在处理
Nx2
形状的数组。因此,我们可以基本上切片第一列和第二列,并分别对它们执行广播减法。好处是我们不会进入
3D
,因此保持内存效率,这也可能转化为性能提升。因此,平方欧几里得距离将被计算如下 -
sq_eucl_dist = (A[:,None,0] - B[:,0])**2 + (A[:,None,1] - B[:,1])**2
让我们计时这三种方法用于计算平方欧几里得距离的时间。
运行时间测试 -
In [75]:
...: A = np.random.rand(200,2)
...: B = np.random.rand(200,2)
...:
In [76]: %timeit ((A[:,None,:] - B[None,:,:])**2).sum(axis=-1)
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