Question

我有一个具有20K行和50列的pandas数据框。我想根据列的欧几里得距离找到此数据框中每行的5个最近邻居。因此，结果是20K * 5的矩阵，其中列是数据帧中最近邻居的ID。

我正在寻找一种解决方案，以尽可能提高效率，最好使用熊猫提供的索引，并行操作或向量化操作。 Scipy KD树非常慢。

有什么主意吗？

Answer 1

对于您的情况，确实似乎Scipy的kd树很慢；查询单个点大约需要80毫秒，我想这将导致整个数据集的总计算时间约为0.08 * 20_000 = 1600s。

高维数据（例如具有50列的数据集）的另一个选项可能是Ball Tree数据结构。如链接中的页面所示：

由于球树节点的球面几何形状，因此它在高维方面的性能要优于KD树，尽管实际性能在很大程度上取决于训练数据的结构。

播放以下代码：

from sklearn.neighbors import NearestNeighbors
import numpy as np

arr = np.random.rand(20_000, 50) * 20
nbrs = NearestNeighbors(n_neighbors = 5, algorithm = 'ball_tree').fit(arr)

%timeit nbrs.kneighbors(arr[:10, :])
# 24.6 ms ± 2.24 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit nbrs.kneighbors(arr[:100, :])
# 209 ms ± 22.7 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit nbrs.kneighbors(arr[:1000, :])
# 2.02 s ± 226 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

看看这些%timeit的结果，似乎该算法大致呈线性扩展，因此对于2万行，您可能会期望它花费大约20_000 / 1_000 * 2 =〜40s。 40秒比您从kd-tree数据结构中最有可能期望的〜1600秒要快得多。

最后，我绝对建议您仔细阅读nearest neighbors页，以使您完全理解它们提供的算法的所有复杂性。

熊猫数据框中所有对的有效k最近邻

1 个答案: