python - 在没有 for 循环的情况下获取到 1D/2D 中每个最近元素的距离

Question

我正在使用 numpy(也可能是 pandas 系列)在 python 中工作，并且正在尝试进行以下计算:

假设我有一个对应于 x 轴上的点的数组:

2, 9, 5, 6, 55, 8

对于这个数组中的每个元素，我想获取到最近元素的距离，因此输出如下所示:

3, 1, 1, 1, 46, 1

我正在尝试找到一种可以缩放到 2D(到最近的 XY 点的距离)的解决方案，并且理想情况下会避免 for 循环。这可能吗？

Answer 1

这里似乎有一个 O(N^2) 解决方案的主题。对于 1D，得到 O(N log N) 非常简单:

x = np.array([2, 9, 5, 6, 55, 8])
i = np.argsort(x)
dist = np.diff(x[i])
min_dist = np.r_[dist[0], np.minimum(dist[1:], dist[:-1]), dist[-1]])
min_dist = min_dist[np.argsort(i)]

这显然不能很好地扩展到多个维度，所以使用 scipy.special.KDTree反而。假设您的数据是 N 维的并且具有形状 (M, N)，您可以这样做

k = KDTree(data)
dist = k.query(data, k=2)[0][:, -1]

Scipy 有一个 KDTree 的 Cython 实现，cKDTree . Sklearn 有一个 sklearn.neighbors.KDTree也有类似的界面。