我最近编写了一个 KDTree,用于在 3-D 空间中进行最近邻搜索,并且在理解 NNS 时遇到了同样的问题,尤其是 wiki 的 3.2。我最终使用了这个算法,它似乎在我的所有测试中都有效:
这是最初的叶子搜索:
public Collection<T> nearestNeighbourSearch(int K, T value) {
if (value==null) return null;
//Map used for results
TreeSet<KdNode> results = new TreeSet<KdNode>(new EuclideanComparator(value));
//Find the closest leaf node
KdNode prev = null;
KdNode node = root;
while (node!=null) {
if (KdNode.compareTo(node.depth, node.k, node.id, value)<0) {
//Greater
prev = node;
node = node.greater;
} else {
//Lesser
prev = node;
node = node.lesser;
}
}
KdNode leaf = prev;
if (leaf!=null) {
//Used to not re-examine nodes
Set<KdNode> examined = new HashSet<KdNode>();
//Go up the tree, looking for better solutions
node = leaf;
while (node!=null) {
//Search node
searchNode(value,node,K,results,examined);
node = node.parent;
}
}
//Load up the collection of the results
Collection<T> collection = new ArrayList<T>(K);
for (KdNode kdNode : results) {
collection.add((T)kdNode.id);
}
return collection;
}
这是从最近的叶节点开始的递归搜索:
private static final <T extends KdTree.XYZPoint> void searchNode(T value, KdNode node, int K, TreeSet<KdNode> results, Set<KdNode> examined) {
examined.add(node);
//Search node
KdNode lastNode = null;
Double lastDistance = Double.MAX_VALUE;
if (results.size()>0) {
lastNode = results.last();
lastDistance = lastNode.id.euclideanDistance(value);
}
Double nodeDistance = node.id.euclideanDistance(value);
if (nodeDistance.compareTo(lastDistance)<0) {
if (results.size()==K && lastNode!=null) results.remove(lastNode);
results.add(node);
} else if (nodeDistance.equals(lastDistance)) {
results.add(node);
} else if (results.size()<K) {
results.add(node);
}
lastNode = results.last();
lastDistance = lastNode.id.euclideanDistance(value);
int axis = node.depth % node.k;
KdNode lesser = node.lesser;
KdNode greater = node.greater;
//Search children branches, if axis aligned distance is less than current distance
if (lesser!=null && !examined.contains(lesser)) {
examined.add(lesser);
double nodePoint = Double.MIN_VALUE;
double valuePlusDistance = Double.MIN_VALUE;
if (axis==X_AXIS) {
nodePoint = node.id.x;
valuePlusDistance = value.x-lastDistance;
} else if (axis==Y_AXIS) {
nodePoint = node.id.y;
valuePlusDistance = value.y-lastDistance;
} else {
nodePoint = node.id.z;
valuePlusDistance = value.z-lastDistance;
}
boolean lineIntersectsCube = ((valuePlusDistance<=nodePoint)?true:false);
//Continue down lesser branch
if (lineIntersectsCube) searchNode(value,lesser,K,results,examined);
}
if (greater!=null && !examined.contains(greater)) {
examined.add(greater);
double nodePoint = Double.MIN_VALUE;
double valuePlusDistance = Double.MIN_VALUE;
if (axis==X_AXIS) {
nodePoint = node.id.x;
valuePlusDistance = value.x+lastDistance;
} else if (axis==Y_AXIS) {
nodePoint = node.id.y;
valuePlusDistance = value.y+lastDistance;
} else {
nodePoint = node.id.z;
valuePlusDistance = value.z+lastDistance;
}
boolean lineIntersectsCube = ((valuePlusDistance>=nodePoint)?true:false);
//Continue down greater branch
if (lineIntersectsCube) searchNode(value,greater,K,results,examined);
}
}
完整的 java 源代码可以在这里找到。