快速受限最近邻搜索

Question

我想为数据矩阵X中的每个点实现一个受限的最近邻搜索。 具体来说，我想在X[mask_i,:]的行中找到X[i,:]的最近邻居，其中mask_i是取决于X[i,:]的逻辑掩码。

下面是使用scipy.spatial.cKDTree class 完成此任务的蛮力方法。 在其中，我将X[i,:]到其最近邻居的距离存储在数组r中的X[mask_i,:]中。

for i in range(n):
    Tree = scipy.spatial.cKDTree(X[mask_i,:]); # mask_i is a logical mask, depending on X[i,:]
    NN = Tree.query(X[i,:], k=2);
    r[i] = NN[0][1];

然而，这种方法在计算上非常昂贵，因为它需要为每个查询训练一个新的 KD-Tree。 我的问题是是否有一种方法可以训练 KD-tree 一次（在 for 循环之前）并在 KD-tree 中查询X[mask_i]中的数据点之间的最近邻居。

如果 KD-trees 不是解决此类问题的正确工具，我将不胜感激任何正确的建议。 非常感谢。