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

Question

我正在使用 numpy（也可能是 pandas 系列）在 python 中工作，并尝试进行以下计算：

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

2, 9, 5, 6, 55, 8

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

3, 1, 1, 1, 46, 1

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

标签： pythonpandasnumpy

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 维的并且具有 shape (M, N)，您可以这样做

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

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