在2D中找到点组之间的最小距离(快速且不太耗费内存)

Question

我有两套点的2D 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

我知道我可以代替cdist一个for循环,保持在只有每个点的最小距离(和索引)A在每一个点B,因为这是我所需要的.我不需要全AxB距离矩阵.但我一直在使用,cdist因为它很快.

有没有办法cdist用一种(差不多？)快的实现替换,但这不会占用那么多内存？

Answer 1

最好的方法是使用专门为最近邻搜索设计的数据结构,例如kd树.例如,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)