I would like to implement a restricted nearest neighbor search for each point in a data matrix X. Specifically, I would like to find the nearest neighbor of X(i,:) among the rows of X(mask_i,:), where mask_i is a logical mask depending on X(i,:).
Below is a brute force way of accomplishing this task, using the KDTreeSearcher model. In it, I store the index of X(i,:) to its nearest neighbor in X(mask_i,:) in an array r.
Mdl = KDTreeSearcher(X(mask_i,:)); # Build KD tree on subset of X satisfying mask_i
r[i] = knnsearch(Mdl, X(i,:)); # Nearest neighbor of X(i,:) in X(mask_i,:)
However, this approach is very computationally expensive since it requires training a new KD-Tree for each query. My question is whether there is a way to train a Kd-tree once (before the for-loop) and query the Kd-tree for nearest neighbors among data points in X(mask_i).
If Kd-trees aren't the right tool for this sort of problem, I would appreciate any suggestions for the right one. Thank you very much.