

 This page uses content from Wikipedia. The original article was at K-d tree. The list of authors can be seen in the page history. As with Rosetta Code, the text of Wikipedia is available under the GNU FDL. (See links for details on variance)

A k-d tree (short for k-dimensional tree) is a space-partitioning data structure for organizing points in a k-dimensional space. k-d trees are a useful data structure for several applications, such as searches involving a multidimensional search key (e.g. range searches and nearest neighbor searches).
k-d trees are a special case of binary space partitioning trees.

k-d trees are not suitable, however, for efficiently finding the nearest neighbor in high dimensional spaces. As a general rule, if the dimensionality is k, the number of points in the data, N, should be N ≫ 2k.
Otherwise, when k-d trees are used with high-dimensional data, most of the points in the tree will be evaluated and the efficiency is no better than exhaustive search, and other methods such as approximate nearest-neighbor are used instead.

Task: Construct a k-d tree and perform a nearest neighbor search for two example data sets:

 The Wikipedia example data of [(2,3), (5,4), (9,6), (4,7), (8,1), (7,2)].
 1000 3-d points uniformly distributed in a 3-d cube.
For the Wikipedia example, find the nearest neighbor to point (9, 2)
For the random data, pick a random location and find the nearest neighbor.

In addition, instrument your code to count the number of nodes visited in the nearest neighbor search. Count a node as visited if any field of it is accessed.

Output should show the point searched for, the point found,
the distance to the point, and the number of nodes visited.

There are variant algorithms for constructing the tree.
You can use a simple median strategy or implement something more efficient.
Variants of the nearest neighbor search include nearest N neighbors, approximate nearest neighbor, and range searches.
You do not have to implement these.
The requirement for this task is specifically the nearest single neighbor.
Also there are algorithms for inserting, deleting, and balancing k-d trees.
These are also not required for the task.

// Implmentation following pseudocode from
// "An introductory tutorial on kd-trees" by Andrew W. Moore,
// Carnegie Mellon University, PDF accessed from:
// http://www.autonlab.org/autonweb/14665
 
import std.typecons, std.math, std.algorithm, std.random, std.range,
       std.traits, core.memory;
 
/// k-dimensional point.
struct Point(size_t k, F) if (isFloatingPoint!F) {
    F[k] data;
    alias data this; // Kills DMD std.algorithm.swap inlining.
                     // Define opIndexAssign and opIndex for dmd.
    enum size_t length = k;
 
    /// Square of the euclidean distance.
    double sqd(in ref Point!(k, F) q) const pure nothrow @nogc {
        double sum = 0;
        foreach (immutable dim, immutable pCoord; data)
            sum += (pCoord - q[dim]) ^^ 2;
        return sum;
    }
}
 
// Following field names in the paper.
// rangeElt would be whatever data is associated with the Point.
// We don't bother with it for this example.
struct KdNode(size_t k, F) {
    Point!(k, F) domElt;
    immutable int split;
    typeof(this)* left, right;
}
 
struct Orthotope(size_t k, F) { /// k-dimensional rectangle.
    Point!(k, F) min, max;
}
 
struct KdTree(size_t k, F) {
    KdNode!(k, F)* n;
    Orthotope!(k, F) bounds;
 
    // Constructs a KdTree from a list of points, also associating the
    // bounds of the tree. The bounds could be computed of course, but
    // in this example we know them already. The algorithm is table
    // 6.3 in the paper.
    this(Point!(k, F)[] pts, in Orthotope!(k, F) bounds_) pure {
        static KdNode!(k, F)* nk2(size_t split)(Point!(k, F)[] exset)
        pure {
            if (exset.empty)
                return null;
            if (exset.length == 1)
                return new KdNode!(k, F)(exset[0], split, null, null);
 
            // Pivot choosing procedure. We find median, then find
            // largest index of points with median value. This
            // satisfies the inequalities of steps 6 and 7 in the
            // algorithm.
            auto m = exset.length / 2;
            topN!((p, q) => p[split] < q[split])(exset, m);
            immutable d = exset[m];
            while (m+1 < exset.length && exset[m+1][split] == d[split])
                m++;
 
            enum nextSplit = (split + 1) % d.length;//cycle coordinates
            return new KdNode!(k, F)(d, split,
                                     nk2!nextSplit(exset[0 .. m]),
                                     nk2!nextSplit(exset[m + 1 .. $]));
        }
 
        this.n = nk2!0(pts);
        this.bounds = bounds_;
    }
}
 
/**
Find nearest neighbor. Return values are:
  nearest neighbor--the ooint within the tree that is nearest p.
  square of the distance to that point.
  a count of the nodes visited in the search.
*/
auto findNearest(size_t k, F)(KdTree!(k, F) t, in Point!(k, F) p)
pure nothrow @nogc {
    // Algorithm is table 6.4 from the paper, with the addition of
    // counting the number nodes visited.
    static Tuple!(Point!(k, F), "nearest",
                  F, "distSqd",
                  int, "nodesVisited")
           nn(KdNode!(k, F)* kd, in Point!(k, F) target,
              Orthotope!(k, F) hr, F maxDistSqd) pure nothrow @nogc {
        if (kd == null)
            return typeof(return)(Point!(k, F)(), F.infinity, 0);
 
        int nodesVisited = 1;
        immutable s = kd.split;
        auto pivot = kd.domElt;
        auto leftHr = hr;
        auto rightHr = hr;
        leftHr.max[s] = pivot[s];
        rightHr.min[s] = pivot[s];
 
        KdNode!(k, F)* nearerKd, furtherKd;
        Orthotope!(k, F) nearerHr, furtherHr;
        if (target[s] <= pivot[s]) {
            //nearerKd, nearerHr = kd.left, leftHr;
            //furtherKd, furtherHr = kd.right, rightHr;
            nearerKd = kd.left;
            nearerHr = leftHr;
            furtherKd = kd.right;
            furtherHr = rightHr;
        } else {
            //nearerKd, nearerHr = kd.right, rightHr;
            //furtherKd, furtherHr = kd.left, leftHr;
            nearerKd = kd.right;
            nearerHr = rightHr;
            furtherKd = kd.left;
            furtherHr = leftHr;
        }
 
        auto n1 = nn(nearerKd, target, nearerHr, maxDistSqd);
        auto nearest = n1.nearest;
        auto distSqd = n1.distSqd;
        nodesVisited += n1.nodesVisited;
 
        if (distSqd < maxDistSqd)
            maxDistSqd = distSqd;
        auto d = (pivot[s] - target[s]) ^^ 2;
        if (d > maxDistSqd)
            return typeof(return)(nearest, distSqd, nodesVisited);
        d = pivot.sqd(target);
        if (d < distSqd) {
            nearest = pivot;
            distSqd = d;
            maxDistSqd = distSqd;
        }
 
        immutable n2 = nn(furtherKd, target, furtherHr, maxDistSqd);
        nodesVisited += n2.nodesVisited;
        if (n2.distSqd < distSqd) {
            nearest = n2.nearest;
            distSqd = n2.distSqd;
        }
 
        return typeof(return)(nearest, distSqd, nodesVisited);
    }
 
    return nn(t.n, p, t.bounds, F.infinity);
}
 
void showNearest(size_t k, F)(in string heading, KdTree!(k, F) kd,
                              in Point!(k, F) p) {
    import std.stdio: writeln;
    writeln(heading, ":");
    writeln("Point:            ", p);
    immutable n = kd.findNearest(p);
    writeln("Nearest neighbor: ", n.nearest);
    writeln("Distance:         ", sqrt(n.distSqd));
    writeln("Nodes visited:    ", n.nodesVisited, "\n");
}
 
void main() {
    static Point!(k, F) randomPoint(size_t k, F)() {
        typeof(return) result;
        foreach (immutable i; 0 .. k)
            result[i] = uniform(F(0), F(1));
        return result;
    }
 
    static Point!(k, F)[] randomPoints(size_t k, F)(in size_t n) {
        return n.iota.map!(_ => randomPoint!(k, F)).array;
    }
 
    import std.stdio, std.conv, std.datetime, std.typetuple;
    rndGen.seed(1); // For repeatable outputs.
 
    alias D2 = TypeTuple!(2, double);
    alias P = Point!D2;
    auto kd1 = KdTree!D2([P([2, 3]), P([5, 4]), P([9, 6]),
                          P([4, 7]), P([8, 1]), P([7, 2])],
                         Orthotope!D2(P([0, 0]), P([10, 10])));
    showNearest("Wikipedia example data", kd1, P([9, 2]));
 
    enum int N = 400_000;
    alias F3 = TypeTuple!(3, float);
    alias Q = Point!F3;
    StopWatch sw;
    GC.disable;
    sw.start;
    auto kd2 = KdTree!F3(randomPoints!F3(N),
                         Orthotope!F3(Q([0, 0, 0]), Q([1, 1, 1])));
    sw.stop;
    GC.enable;
    showNearest(text("k-d tree with ", N,
                     " random 3D ", F3[1].stringof,
                     " points (construction time: ",
                     sw.peek.msecs, " ms)"), kd2, randomPoint!F3);
 
    sw.reset;
    sw.start;
    enum int M = 10_000;
    size_t visited = 0;
    foreach (immutable _; 0 .. M) {
        immutable n = kd2.findNearest(randomPoint!F3);
        visited += n.nodesVisited;
    }
    sw.stop;
 
    writefln("Visited an average of %0.2f nodes on %d searches " ~
             "in %d ms.", visited / double(M), M, sw.peek.msecs);
}