python - 查找二维点组之间的最小距离(快速且不太消耗内存)

Question

我在二维 A 和 B 中有两组点，我需要找到 A 中每个点的最小距离，以B 中的一个点。到目前为止，我一直在使用 SciPy 的 cdist使用下面的代码

import numpy as np
from scipy.spatial.distance import cdist

def ABdist(A, B):
    # Distance to all points in B, for each point in A.
    dist = cdist(A, B, 'euclidean')
    # Indexes to minimum distances.
    min_dist_idx = np.argmin(dist, axis=1)
    # Store only the minimum distances for each point in A, to a point in B.
    min_dists = [dist[i][md_idx] for i, md_idx in enumerate(min_dist_idx)]

    return min_dist_idx, min_dists

N = 10000
A = np.random.uniform(0., 5000., (N, 2))
B = np.random.uniform(0., 5000., (N, 2))

min_dist_idx, min_dists = ABdist(A, B)

对于 N 的小值来说效果很好。但是现在集合的长度从 N=10000 增加到 N=35000 我遇到了一个

    dm = np.zeros((mA, mB), dtype=np.double)
MemoryError

我知道我可以用 for 循环替换 cdist，它只保留 A 中每个点到 中每个点的最小距离(和索引) B，因为这就是我所需要的。我不需要完整的 AxB 距离矩阵。但我一直在使用 cdist 正是因为它很快。

有没有一种方法可以用(几乎？)一样快但不会占用那么多内存的实现来替换 cdist？

Answer 1

最好的方法是使用专门为最近邻搜索设计的数据结构，例如 k-d tree .例如，SciPy 的 cKDTree允许您以这种方式解决问题:

from scipy.spatial import cKDTree
min_dists, min_dist_idx = cKDTree(B).query(A, 1)

无论是在计算还是内存使用方面，结果都比任何基于广播的方法更有效。

例如，即使有 1,000,000 个点，计算也不会耗尽内存，并且在我的笔记本电脑上只需要几秒钟:

N = 1000000
A = np.random.uniform(0., 5000., (N, 2))
B = np.random.uniform(0., 5000., (N, 2))

%timeit cKDTree(B).query(A, 1)
# 3.25 s ± 17.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)