Task: K-d tree




Go: 


package main

import (
    "fmt"
    "math"
    "math/rand"
    "sort"
    "time"
)


type point []float64


func (p point) sqd(q point) float64 {
    var sum float64
    for dim, pCoord := range p {
        d := pCoord - q[dim]
        sum += d * d
    }
    return sum
}




type kdNode struct {
    domElt      point
    split       int
    left, right *kdNode
}   

type kdTree struct {
    n      *kdNode
    bounds hyperRect
}
    
type hyperRect struct {
    min, max point
}



func (hr hyperRect) copy() hyperRect {
    return hyperRect{append(point{}, hr.min...), append(point{}, hr.max...)}
}   
    



func newKd(pts []point, bounds hyperRect) kdTree {
    var nk2 func([]point, int) *kdNode
    nk2 = func(exset []point, split int) *kdNode {
        if len(exset) == 0 {
            return nil
        }
        
        
        
        sort.Sort(part{exset, split})
        m := len(exset) / 2
        d := exset[m]
        for m+1 < len(exset) && exset[m+1][split] == d[split] {
            m++
        }
        
        s2 := split + 1
        if s2 == len(d) {
            s2 = 0
        }
        return &kdNode{d, split, nk2(exset[:m], s2), nk2(exset[m+1:], s2)}
    }
    return kdTree{nk2(pts, 0), bounds}
}



type part struct {
    pts   []point
    dPart int
}


func (p part) Len() int { return len(p.pts) }
func (p part) Less(i, j int) bool {
    return p.pts[i][p.dPart] < p.pts[j][p.dPart]
}
func (p part) Swap(i, j int) { p.pts[i], p.pts[j] = p.pts[j], p.pts[i] }





func (t kdTree) nearest(p point) (best point, bestSqd float64, nv int) {
    return nn(t.n, p, t.bounds, math.Inf(1))
}



func nn(kd *kdNode, target point, hr hyperRect,
    maxDistSqd float64) (nearest point, distSqd float64, nodesVisited int) {
    if kd == nil {
        return nil, math.Inf(1), 0
    }
    nodesVisited++
    s := kd.split
    pivot := kd.domElt
    leftHr := hr.copy()
    rightHr := hr.copy()
    leftHr.max[s] = pivot[s]
    rightHr.min[s] = pivot[s]
    targetInLeft := target[s] <= pivot[s]
    var nearerKd, furtherKd *kdNode
    var nearerHr, furtherHr hyperRect
    if targetInLeft {
        nearerKd, nearerHr = kd.left, leftHr
        furtherKd, furtherHr = kd.right, rightHr
    } else {
        nearerKd, nearerHr = kd.right, rightHr
        furtherKd, furtherHr = kd.left, leftHr
    }
    var nv int
    nearest, distSqd, nv = nn(nearerKd, target, nearerHr, maxDistSqd)
    nodesVisited += nv
    if distSqd < maxDistSqd {
        maxDistSqd = distSqd
    }
    d := pivot[s] - target[s]
    d *= d
    if d > maxDistSqd {
        return
    }
    if d = pivot.sqd(target); d < distSqd {
        nearest = pivot
        distSqd = d
        maxDistSqd = distSqd
    }
    tempNearest, tempSqd, nv := nn(furtherKd, target, furtherHr, maxDistSqd)
    nodesVisited += nv
    if tempSqd < distSqd {
        nearest = tempNearest
        distSqd = tempSqd
    }
    return
}

func main() {
    rand.Seed(time.Now().Unix())
    kd := newKd([]point{{2, 3}, {5, 4}, {9, 6}, {4, 7}, {8, 1}, {7, 2}},
        hyperRect{point{0, 0}, point{10, 10}})
    showNearest("WP example data", kd, point{9, 2})
    kd = newKd(randomPts(3, 1000), hyperRect{point{0, 0, 0}, point{1, 1, 1}})
    showNearest("1000 random 3d points", kd, randomPt(3))
}   
    
func randomPt(dim int) point {
    p := make(point, dim)
    for d := range p {
        p[d] = rand.Float64()
    }
    return p
}   
    
func randomPts(dim, n int) []point {
    p := make([]point, n)
    for i := range p {
        p[i] = randomPt(dim) 
    } 
    return p
}
    
func showNearest(heading string, kd kdTree, p point) {
    fmt.Println()
    fmt.Println(heading)
    fmt.Println("point:           ", p)
    nn, ssq, nv := kd.nearest(p)
    fmt.Println("nearest neighbor:", nn)
    fmt.Println("distance:        ", math.Sqrt(ssq))
    fmt.Println("nodes visited:   ", nv)
}


Java: import java.util.*;

public class KdTree {
    private int dimensions_;
    private Node root_ = null;
    private Node best_ = null;
    private double bestDistance_ = 0;
    private int visited_ = 0;
    
    public KdTree(int dimensions, List<Node> nodes) {
        dimensions_ = dimensions;
        root_ = makeTree(nodes, 0, nodes.size(), 0);
    }
    
    public Node findNearest(Node target) {
        if (root_ == null)
            throw new IllegalStateException("Tree is empty!");
        best_ = null;
        visited_ = 0;
        bestDistance_ = 0;
        nearest(root_, target, 0);
        return best_;
    }
    
    public int visited() {
        return visited_;
    }
    
    public double distance() {
        return Math.sqrt(bestDistance_);
    }
    
    private void nearest(Node root, Node target, int index) {
        if (root == null)
            return;
        ++visited_;
        double d = root.distance(target);
        if (best_ == null || d < bestDistance_) {
            bestDistance_ = d;
            best_ = root;
        }
        if (bestDistance_ == 0)
            return;
        double dx = root.get(index) - target.get(index);
        index = (index + 1) % dimensions_;
        nearest(dx > 0 ? root.left_ : root.right_, target, index);
        if (dx * dx >= bestDistance_)
            return;
        nearest(dx > 0 ? root.right_ : root.left_, target, index);
    }
    
    private Node makeTree(List<Node> nodes, int begin, int end, int index) {
        if (end <= begin)
            return null;
        int n = begin + (end - begin)/2;
        Node node = QuickSelect.select(nodes, begin, end - 1, n, new NodeComparator(index));
        index = (index + 1) % dimensions_;
        node.left_ = makeTree(nodes, begin, n, index);
        node.right_ = makeTree(nodes, n + 1, end, index);
        return node;
    }
    
    private static class NodeComparator implements Comparator<Node> {
        private int index_;

        private NodeComparator(int index) {
            index_ = index;
        }
        public int compare(Node n1, Node n2) {
            return Double.compare(n1.get(index_), n2.get(index_));
        }
    }
    
    public static class Node {
        private double[] coords_;
        private Node left_ = null;
        private Node right_ = null;

        public Node(double[] coords) {
            coords_ = coords;
        }
        public Node(double x, double y) {
            this(new double[]{x, y});
        }
        public Node(double x, double y, double z) {
            this(new double[]{x, y, z});
        }
        double get(int index) {
            return coords_[index];
        }
        double distance(Node node) {
            double dist = 0;
            for (int i = 0; i < coords_.length; ++i) {
                double d = coords_[i] - node.coords_[i];
                dist += d * d;
            }
            return dist;
        }
        public String toString() {
            StringBuilder s = new StringBuilder("(");
            for (int i = 0; i < coords_.length; ++i) {
                if (i > 0)
                    s.append(", ");
                s.append(coords_[i]);
            }
            s.append(')');
            return s.toString();
        }
    }
}