使用Numpy高效地计算欧几里得距离矩阵

31

我有一组二维空间中的点,需要计算每个点到其他所有点的距离。

我的点数量相对较少,最多可能只有100个。但由于我需要经常快速地进行计算以确定这些移动点之间的关系,并且我知道遍历这些点可能会导致O(n^2)复杂度的问题,因此我正在寻找利用numpy矩阵运算(或scipy)的方法。

目前在我的代码中,每个对象的坐标存储在其类中。然而,当我更新类坐标时,我也可以将它们更新到numpy数组中。

class Cell(object):
    """Represents one object in the field."""
    def __init__(self,id,x=0,y=0):
        self.m_id = id
        self.m_x = x
        self.m_y = y

我想创建一个欧几里得距离矩阵来避免重复,但也许您有更聪明的数据结构。

我很乐意接受关于巧妙算法的指针。

此外,我注意到有类似的问题涉及欧几里得距离和numpy,但没有找到直接解决高效填充完整距离矩阵的答案。


3
这可能会有所帮助:scipy.spatial.distance.pdist 是一个用于计算欧几里得距离和其他距离度量的函数。 - Carsten
4
无论如何,复杂度都将是O(n ^ 2):针对一般点集,您可以做的最好的事情是仅计算“ n *(n-1)/ 2”个距离,这仍然是O(n ^ 2)。 - Jaime
如果可以使用scipy,请考虑使用scipy.spatial.distance_matrix - Eric Gopak
5个回答

42

您可以利用 complex 类型:

# build a complex array of your cells
z = np.array([complex(c.m_x, c.m_y) for c in cells])

第一种解决方案

# mesh this array so that you will have all combinations
m, n = np.meshgrid(z, z)
# get the distance via the norm
out = abs(m-n)

第二种解决方案

网格生成是主要思路。但是numpy很聪明,因此您不必生成mn。只需使用z的转置版本计算差异即可。网格会自动完成:

out = abs(z[..., np.newaxis] - z)

第三解决方案

如果直接将z设置为一个二维数组,您可以使用z.T而不是奇怪的z[..., np.newaxis]。所以最终,您的代码将如下所示:

z = np.array([[complex(c.m_x, c.m_y) for c in cells]]) # notice the [[ ... ]]
out = abs(z.T-z)

例子

>>> z = np.array([[0.+0.j, 2.+1.j, -1.+4.j]])
>>> abs(z.T-z)
array([[ 0.        ,  2.23606798,  4.12310563],
       [ 2.23606798,  0.        ,  4.24264069],
       [ 4.12310563,  4.24264069,  0.        ]])

作为补充,您可能希望之后删除重复项,只选择上三角:

>>> np.triu(out)
array([[ 0.        ,  2.23606798,  4.12310563],
       [ 0.        ,  0.        ,  4.24264069],
       [ 0.        ,  0.        ,  0.        ]])

一些基准测试

>>> timeit.timeit('abs(z.T-z)', setup='import numpy as np;z = np.array([[0.+0.j, 2.+1.j, -1.+4.j]])')
4.645645342274779
>>> timeit.timeit('abs(z[..., np.newaxis] - z)', setup='import numpy as np;z = np.array([0.+0.j, 2.+1.j, -1.+4.j])')
5.049334864854522
>>> timeit.timeit('m, n = np.meshgrid(z, z); abs(m-n)', setup='import numpy as np;z = np.array([0.+0.j, 2.+1.j, -1.+4.j])')
22.489568296184686

8
你找到那段距离了吗?如果是的话,你把我丢在哪里了?那是在哪里发生的? - Wes Modes
@WesModes,回答有些晚了,但仍可能有用。复数基本上是一个二维点。两个复数的差是一个复数。复数的绝对值是从(0,0)到该点的距离。 - newtover

14
如果你不需要完整的距离矩阵,使用kd树会更好。考虑使用scipy.spatial.cKDTree或者sklearn.neighbors.KDTree。这是因为kd树可以在O(n log n)时间内找到k个最近的邻居,从而避免了计算所有n乘以n距离的O(n**2)复杂度。

13

Jake Vanderplas在《Python数据科学手册》中给出了使用广播的示例,这与@shx2提出的非常相似。

import numpy as np
rand = random.RandomState(42)
X = rand.rand(3, 2)  
dist_sq = np.sum((X[:, np.newaxis, :] - X[np.newaxis, :, :]) ** 2, axis = -1)

dist_sq
array([[0.        , 0.18543317, 0.81602495],
       [0.18543317, 0.        , 0.22819282],
       [0.81602495, 0.22819282, 0.        ]])

scipy.spatial.distance.cdist 的速度比这个快,经过我的测试是9倍 - Tweakimp
2
@Tweakimp - 你应该写一个带有%timeit调用的答案,也许是针对一个小的(10x10)和大的(1,000,000 x 1,000,000)距离矩阵。这将为人们提供非常有用的信息! - Rich Pauloo
我无法在我的Jupyter笔记本中使用%timeit,因为我使用的是在线变体,对于如此大的数组它会耗尽内存。 - Tweakimp
这是一个超级快速的解决方案。 - Ramin Melikov
这个解决方案是广播的一个很好的例子,但它消耗了 Θ(n^2 * d) 的内存(其中 n 是向量的数量,d 是维度),而最优解只会消耗 O(n^2)。 (由 /usr/bin/time -v 确认。) - japreiss
你会如何计算曼哈顿距离? - Ramin Melikov

8

以下是使用numpy的方法:

import numpy as np

x = np.array([0,1,2])
y = np.array([2,4,6])

# take advantage of broadcasting, to make a 2dim array of diffs
dx = x[..., np.newaxis] - x[np.newaxis, ...]
dy = y[..., np.newaxis] - y[np.newaxis, ...]
dx
=> array([[ 0, -1, -2],
          [ 1,  0, -1],
          [ 2,  1,  0]])

# stack in one array, to speed up calculations
d = np.array([dx,dy])
d.shape
=> (2, 3, 3)

现在只需要沿着0轴计算L2范数(如此处所讨论:这里)。
(d**2).sum(axis=0)**0.5
=> array([[ 0.        ,  2.23606798,  4.47213595],
          [ 2.23606798,  0.        ,  2.23606798],
          [ 4.47213595,  2.23606798,  0.        ]])

如果x或y很大,这实际上需要相当多的内存,同时速度也很慢。SciPy的距离矩阵应该会快得多。 - Lars Gebraad

4
如果您正在寻找最高效的计算方式,请使用SciPy的cdist()(或pdist(),如果您只需要成对距离向量而不是完整的距离矩阵),如Tweakimp的评论所建议。正如他所说,它比基于向量化和广播提出的方法由RichPauloo和shx2要快得多。原因是SciPy的cdist()pdist()在底层使用for循环和C实现进行度量计算,这甚至比向量化更快。
顺便提一下,如果你能使用SciPy并且仍然喜欢使用广播方法,那么你不必自己实现它,因为distance_matrix()函数是纯Python实现的,利用了广播和向量化(源代码文档)。
值得一提的是,cdist()/pdist()在内存方面也比广播更有效率,因为它逐个计算距离,避免创建n*n*d元素的数组,其中n是点数,d是点的维度。
实验
我进行了一些简单的实验,比较了SciPy的cdist()distance_matrix()和NumPy中的广播实现的性能。我使用Python的时间模块中的perf_counter_ns()来测量时间,并且所有结果都是在2D空间中使用np.float64数据类型的10000个点上进行的10次平均运行(在Python 3.8.10,Windows 10下测试,使用Ryzen 2700和16 GB RAM):

  • cdist() - 0.6724秒
  • distance_matrix() - 3.0128秒
  • 我的NumPy实现 - 3.6931秒

如果有人想要重现实验,请参考以下代码:

from scipy.spatial import *
import numpy as np
from time import perf_counter_ns


def dist_mat_custom(a, b):
    return np.sqrt(np.sum(np.square(a[:, np.newaxis, :] - b[np.newaxis, :, :]), axis=-1))


results = []
size = 10000
it_num = 10
for i in range(it_num):
    a = np.random.normal(size=(size, 2))
    b = np.random.normal(size=(size, 2))
    start = perf_counter_ns()
    c = distance_matrix(a, b)
    #c = dist_mat_custom(a, b)
    #c = distance.cdist(a, b)
    results.append(perf_counter_ns() - start)
print(np.mean(results) / 1e9)

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接