63 lines
2.4 KiB
Text
63 lines
2.4 KiB
Text
class KdTree{
|
|
const NEAREST=0, DIST_SQD=1, NODES_VISITED=2;
|
|
class KdNode{
|
|
fcn init(_dom_elt,_split,_left,_right){
|
|
var dom_elt=_dom_elt.apply("toFloat"),
|
|
split=_split, left=_left, right=_right;
|
|
}
|
|
}
|
|
fcn init(pts,_bounds){ // pts is points is ( (x,y,..),..)
|
|
var n=fcn(split, exset){
|
|
if(not exset) return(Void);
|
|
exset=exset.copy() // mutable lists sort much quicker than read only
|
|
.sort('wrap(abc,xyz){ abc[split]<xyz[split] });
|
|
m,d:=exset.len()/2, exset[m];
|
|
while(m+1<exset.len() and exset[m + 1][split]==d[split]){ m+=1 }
|
|
s2:=(split+1)%d.len(); # cycle coordinates
|
|
KdNode(d,split,self.fcn(s2,exset[0,m]),
|
|
self.fcn(s2,exset[m+1,*]))
|
|
}(0,pts);
|
|
var bounds=_bounds;
|
|
}
|
|
fcn findNearest(k,p){
|
|
fcn(node,target,hr,max_dist_sqd,k){
|
|
if(not node) return(k.pump(List,0.0), (0.0).MAX, 0); // fake node far away
|
|
nodes_visited,s,pivot:=1, node.split, node.dom_elt;
|
|
left_hr,right_hr:=hr.copy(True), hr.copy(True); // deep-ish copy
|
|
left_hr.max[s]=right_hr.min[s]=pivot[s];
|
|
|
|
reg nearer_node, nearer_hr, further_node, further_hr;
|
|
if(target[s]<=pivot[s]){
|
|
nearer_node, nearer_hr = node.left, left_hr;
|
|
further_node, further_hr = node.right, right_hr;
|
|
}else{
|
|
nearer_node, nearer_hr = node.right, right_hr;
|
|
further_node, further_hr = node.left, left_hr;
|
|
}
|
|
n1:=self.fcn(nearer_node, target, nearer_hr, max_dist_sqd, k);
|
|
nearest,dist_sqd:=n1; // n1 is ( (a,b),distance,#visited )
|
|
nodes_visited+=n1[NODES_VISITED];
|
|
|
|
if(dist_sqd<max_dist_sqd) max_dist_sqd=dist_sqd;
|
|
d:=(pivot[s] - target[s]).pow(2);
|
|
if(d>max_dist_sqd) return(nearest, dist_sqd, nodes_visited);
|
|
d:=sqd(pivot,target);
|
|
if(d<dist_sqd) nearest,dist_sqd,max_dist_sqd=pivot,d,dist_sqd;
|
|
n2:=self.fcn(further_node, target, further_hr, max_dist_sqd,k);
|
|
nodes_visited+=n2[NODES_VISITED];
|
|
if(n2[DIST_SQD]<dist_sqd) nearest,dist_sqd=n2;
|
|
return(nearest, dist_sqd, nodes_visited)
|
|
}(n, p, bounds, (0.0).MAX,k)
|
|
}
|
|
fcn [private] sqd(p1,p2){ // point deltas squared and summed
|
|
p1.zipWith(fcn(a,b){ (a-b).pow(2) }, p2).sum(0.0)
|
|
}
|
|
fcn show_nearest(k, heading, p){
|
|
println(heading,":");
|
|
println("Point: ",p);
|
|
n:=findNearest(k,p);
|
|
println("Nearest neighbor:", n[NEAREST]);
|
|
println("Distance: ", n[DIST_SQD].sqrt());
|
|
println("Nodes visited: ", n[NODES_VISITED], "\n");
|
|
}
|
|
}
|