158 lines
4.1 KiB
D
158 lines
4.1 KiB
D
import std.stdio, std.algorithm, std.math, std.random;
|
|
|
|
struct KdNode(size_t dim) {
|
|
double[dim] x;
|
|
KdNode* left, right;
|
|
}
|
|
|
|
// See QuickSelect method.
|
|
KdNode!dim* findMedian(size_t idx, size_t dim)(KdNode!dim[] nodes) pure nothrow @nogc {
|
|
auto start = nodes.ptr;
|
|
auto end = &nodes[$ - 1] + 1;
|
|
|
|
if (end <= start)
|
|
return null;
|
|
if (end == start + 1)
|
|
return start;
|
|
|
|
auto md = start + (end - start) / 2;
|
|
|
|
while (true) {
|
|
immutable double pivot = md.x[idx];
|
|
|
|
swap(md.x, (end - 1).x); // Swaps the whole arrays x.
|
|
auto store = start;
|
|
foreach (p; start .. end) {
|
|
if (p.x[idx] < pivot) {
|
|
if (p != store)
|
|
swap(p.x, store.x);
|
|
store++;
|
|
}
|
|
}
|
|
swap(store.x, (end - 1).x);
|
|
|
|
// Median has duplicate values.
|
|
if (store.x[idx] == md.x[idx])
|
|
return md;
|
|
|
|
if (store > md)
|
|
end = store;
|
|
else
|
|
start = store;
|
|
}
|
|
}
|
|
|
|
KdNode!dim* makeTree(size_t dim, size_t i = 0)(KdNode!dim[] nodes)
|
|
pure nothrow @nogc {
|
|
if (!nodes.length)
|
|
return null;
|
|
|
|
auto n = nodes.findMedian!i;
|
|
if (n != null) {
|
|
enum i2 = (i + 1) % dim;
|
|
immutable size_t nPos = n - nodes.ptr;
|
|
n.left = makeTree!(dim, i2)(nodes[0 .. nPos]);
|
|
n.right = makeTree!(dim, i2)(nodes[nPos + 1 .. $]);
|
|
}
|
|
|
|
return n;
|
|
}
|
|
|
|
void nearest(size_t dim)(in KdNode!dim* root,
|
|
in ref KdNode!dim nd,
|
|
in size_t i,
|
|
ref const(KdNode!dim)* best,
|
|
ref double bestDist,
|
|
ref size_t nVisited) pure nothrow @safe @nogc {
|
|
static double dist(in ref KdNode!dim a, in ref KdNode!dim b)
|
|
pure nothrow @nogc {
|
|
double result = 0;
|
|
static foreach (i; 0 .. dim)
|
|
result += (a.x[i] - b.x[i]) ^^ 2;
|
|
return result;
|
|
}
|
|
|
|
if (root == null)
|
|
return;
|
|
|
|
immutable double d = dist(*root, nd);
|
|
immutable double dx = root.x[i] - nd.x[i];
|
|
immutable double dx2 = dx ^^ 2;
|
|
nVisited++;
|
|
|
|
if (!best || d < bestDist) {
|
|
bestDist = d;
|
|
best = root;
|
|
}
|
|
|
|
// If chance of exact match is high.
|
|
if (!bestDist)
|
|
return;
|
|
|
|
immutable i2 = (i + 1 >= dim) ? 0 : i + 1;
|
|
|
|
nearest!dim(dx > 0 ? root.left : root.right,
|
|
nd, i2, best, bestDist, nVisited);
|
|
if (dx2 >= bestDist)
|
|
return;
|
|
nearest!dim(dx > 0 ? root.right : root.left,
|
|
nd, i2, best, bestDist, nVisited);
|
|
}
|
|
|
|
void randPt(size_t dim)(ref KdNode!dim v, ref Xorshift rng)
|
|
pure nothrow @safe @nogc {
|
|
static foreach (i; 0 .. dim)
|
|
v.x[i] = rng.uniform01;
|
|
}
|
|
|
|
/// smallTest
|
|
unittest {
|
|
KdNode!2[] wp = [{[2, 3]}, {[5, 4]}, {[9, 6]},
|
|
{[4, 7]}, {[8, 1]}, {[7, 2]}];
|
|
KdNode!2 thisPt = {[9, 2]};
|
|
|
|
auto root = makeTree(wp);
|
|
|
|
const(KdNode!2)* found = null;
|
|
double bestDist = 0;
|
|
size_t nVisited = 0;
|
|
root.nearest(thisPt, 0, found, bestDist, nVisited);
|
|
|
|
writefln("WP tree:\n Searching for %s\n" ~
|
|
" Found %s, dist = %g\n Seen %d nodes.\n",
|
|
thisPt.x, found.x, sqrt(bestDist), nVisited);
|
|
}
|
|
|
|
/// bigTest
|
|
unittest {
|
|
enum N = 1_000_000;
|
|
enum testRuns = 100_000;
|
|
|
|
auto bigTree = new KdNode!3[N];
|
|
auto rng = 1.Xorshift;
|
|
foreach (ref node; bigTree)
|
|
randPt(node, rng);
|
|
|
|
auto root = makeTree(bigTree);
|
|
KdNode!3 thisPt;
|
|
randPt(thisPt, rng);
|
|
|
|
const(KdNode!3)* found = null;
|
|
double bestDist = 0;
|
|
size_t nVisited = 0;
|
|
root.nearest(thisPt, 0, found, bestDist, nVisited);
|
|
|
|
writefln("Big tree (%d nodes):\n Searching for %s\n" ~ " Found %s, dist = %g\n Seen %d nodes.", N, thisPt.x, found.x, sqrt(bestDist), nVisited);
|
|
|
|
size_t sum = 0;
|
|
foreach (immutable _; 0 .. testRuns) {
|
|
found = null;
|
|
nVisited = 0;
|
|
randPt(thisPt, rng);
|
|
nearest!3(root, thisPt, 0, found, bestDist, nVisited);
|
|
sum += nVisited;
|
|
}
|
|
writefln("\nBig tree:\n Visited %d nodes for %d random " ~
|
|
"searches (%.2f per lookup).",
|
|
sum, testRuns, sum / double(testRuns));
|
|
}
|