在 Python 中，需要一种将 kdtree 索引映射到值的有效方法

Question

我正在使用来自 scikit-learn 的 kdtree 和一个非常大的数据集。

我可以让 kdtree 在合理的时间内（在我的机器上 20 分钟）执行查询，但我无法在不到 1 小时的任何时间内将索引映射到它们所代表的值（我在 1 小时后停止等待)。

我加载了 2 个 csv 文件（train.csv 有 29M 记录，test.csv 有 8M 记录）。我对 3 把钥匙感兴趣。 'x','y' 是浮点数，'placeid' 是字符串。

from sklearn.neighbors import KDTree
import pandas as pd

train = pd.read_csv("train.csv")
test = pd.read.csv("test.csv")

tree = KDTree(train[['x','y']])
_, indexes = tree.query(test[['x','y']],k=30)

# takes 20 minutes to get here.  Here is the code that takes more than an hour

result = [[train.iloc[idx].place_id for idx in idx_set] for idx_set in indexes]

有没有更快的方法来做到这一点？我的目标是将所有从 KDTree 返回的索引映射到 place_ids。

Answer 1

也许你可以试一试，因为你不想与query保持距离：

indexes = tree.query(test[['x','y']],k=30,return_distance=False,dualtree=True,sort_results=False)

这可能会减少第一部分的一些计算时间。 对于第二部分，我正在考虑扁平化或reshape 索引和切片place_id，而不是双循环。可以提供result的格式吗？它只是一个简单的列表吗？