使用欧几里得距离在numpy数组列表中找到最近邻居的numpy数组。

Question

我有一个n维向量，希望在一个包含许多n维向量的列表中，使用欧几里得距离找到其k个最近邻。

我编写了下面的代码（其中k=10），虽然可以工作，但运行速度太慢，想知道是否有更优化的解决方案。

def nearest_neighbors(value, array, nbr_neighbors=1):
    return np.argsort(np.array([np.linalg.norm(value-x) for x in array]))[:nbr_neighbors]

- Bastien Beurier

Answer 1

使用scipy的kd-tree。

这里有一个小例子可供参考。

许多人似乎抱怨性能，并推荐sklearn的实现，尽管它使用了数据结构（内部使用sklearn.neighbors）！

- sascha

---

SciPy拥有cKDTree，它比KDTree快得多。 - teekarna

更新链接：cKDtree。 - Matti Wens

在 SciPy 中，除了向后兼容性外，使用 cKDtree 而不是 KDtree 已不再推荐，因为它们现在是相同的。从上面链接的 cKDtree 文档页面中：“在 SciPy v1.6.0 之前，cKDTree 的性能更好，功能略有不同，但现在这两个名称仅出于向后兼容性的原因而存在。如果与 SciPy <1.6 兼容性不是问题，请优先选择 KDTree。” - realityChemist

Answer 2

正如Sascha所说，我最终使用了scipy库（但使用了NearestNeighbors方法），将计算时间从50小时缩短到了36分钟。这种计算我不应该尝试自己重新实现，因为专用库对此进行了更多的优化。 NearestNeighbors方法还允许您传入值列表，并返回每个值的k个最近邻居。

最终代码如下：

def nearest_neighbors(values, all_values, nbr_neighbors=10):
    nn = NearestNeighbors(nbr_neighbors, metric='cosine', algorithm='brute').fit(all_values)
    dists, idxs = nn.kneighbors(values)

- Bastien Beurier

---

你应该选择他的答案作为最佳答案，并将其作为最终编辑添加到你的帖子中。 - dawg

这仅适用于二维数据。更高维度将无法工作。 - FindOutIslamNow

Answer 3

我建议使用Scipy的pdist函数通过暴力方法找到成对距离：https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.distance.pdist.html。由于pdist高度优化，因此速度应该非常快。然后对于每个元素选择k个最近的。

- Zhenlei Cai