@thi.ng/geom-accel
Version:
n-D spatial indexing data structures with a shared ES6 Map/Set-like API
288 lines (287 loc) • 7.22 kB
JavaScript
import { ensureArray } from "@thi.ng/arrays/ensure-array";
import { Heap } from "@thi.ng/heaps/heap";
import { EPS } from "@thi.ng/math/api";
import { map } from "@thi.ng/transducers/map";
import { distSq } from "@thi.ng/vectors/distsq";
import { CMP, __addResults, __into } from "./utils.js";
class KdNode {
d;
parent;
l;
r;
k;
v;
constructor(parent, dim, key, val) {
this.parent = parent;
this.d = dim;
this.k = key;
this.v = val;
}
get height() {
return 1 + Math.max(this.l ? this.l.height : 0, this.r ? this.r.height : 0);
}
}
class KdTreeMap {
constructor(dim, pairs, distanceFn = distSq) {
this.distanceFn = distanceFn;
this.dim = dim;
this._size = 0;
this.root = pairs ? this.buildTree(ensureArray(pairs), 0) : void 0;
}
dim;
root;
_size;
*[Symbol.iterator]() {
let queue = this.root ? [this.root] : [];
while (queue.length) {
const n = queue.pop();
if (n) {
yield [n.k, n.v];
queue.push(n.r, n.l);
}
}
}
*keys() {
let queue = this.root ? [this.root] : [];
while (queue.length) {
const n = queue.pop();
if (n) {
yield n.k;
queue.push(n.r, n.l);
}
}
}
values() {
return map((p) => p[1], this);
}
get size() {
return this._size;
}
get height() {
return this.root ? this.root.height : 0;
}
get ratio() {
return this._size ? this.height / Math.log2(this._size) : 0;
}
copy() {
return new KdTreeMap(this.dim, this, this.distanceFn);
}
clear() {
delete this.root;
this._size = 0;
}
empty() {
return new KdTreeMap(this.dim, void 0, this.distanceFn);
}
set(key, val, eps = EPS) {
eps = Math.max(0, eps);
eps *= eps;
const search = (node, parent2) => node ? search(key[node.d] < node.k[node.d] ? node.l : node.r, node) : parent2;
let parent;
if (this.root) {
parent = __nearest1(
key,
[eps, void 0],
this.dim,
this.root,
this.distanceFn
)[1];
if (parent) {
parent.v = val;
return false;
}
parent = search(this.root, void 0);
const dim = parent.d;
parent[key[dim] < parent.k[dim] ? "l" : "r"] = new KdNode(
parent,
(dim + 1) % this.dim,
key,
val
);
} else {
this.root = new KdNode(void 0, 0, key, val);
}
this._size++;
return true;
}
into(pairs, eps = EPS) {
return __into(this, pairs, eps);
}
remove(key) {
const node = __find(key, this.root, 0);
if (node) {
__remove(node) && (this.root = void 0);
this._size--;
return true;
}
return false;
}
has(key, eps = EPS) {
return !!this.root && !!__nearest1(
key,
[eps * eps, void 0],
this.dim,
this.root,
this.distanceFn
)[1];
}
get(key, eps = EPS) {
if (this.root) {
const node = __nearest1(
key,
[eps * eps, void 0],
this.dim,
this.root,
this.distanceFn
)[1];
return node ? node.v : void 0;
}
}
query(q, maxDist, limit, acc) {
return this.doSelect(q, (x) => [x.k, x.v], maxDist, limit, acc);
}
queryKeys(q, maxDist, limit, acc) {
return this.doSelect(q, (x) => x.k, maxDist, limit, acc);
}
queryValues(q, maxDist, limit, acc) {
return this.doSelect(q, (x) => x.v, maxDist, limit, acc);
}
doSelect(q, f, maxDist, maxNum = 1, acc = []) {
if (!this.root) return [];
maxDist *= maxDist;
if (maxNum === 1) {
const sel = __nearest1(
q,
[maxDist, void 0],
this.dim,
this.root,
this.distanceFn
)[1];
sel && acc.push(f(sel));
} else {
const nodes = new Heap(
[[maxDist, void 0]],
{
compare: CMP
}
);
__nearest(q, nodes, this.dim, maxNum, this.root, this.distanceFn);
return __addResults(f, nodes.values, acc);
}
return acc;
}
buildTree(points, depth, parent) {
const n = points.length;
if (n === 0) {
return;
}
this._size++;
let dim = depth % this.dim;
if (n === 1) {
return new KdNode(parent, dim, ...points[0]);
}
points.sort((a, b) => a[0][dim] - b[0][dim]);
const med = n >>> 1;
const node = new KdNode(parent, dim, ...points[med]);
node.l = this.buildTree(points.slice(0, med), depth + 1, node);
node.r = this.buildTree(points.slice(med + 1), depth + 1, node);
return node;
}
}
const __find = (p, node, epsSq) => {
if (!node) return;
return distSq(p, node.k) <= epsSq ? node : __find(p, p[node.d] < node.k[node.d] ? node.l : node.r, epsSq);
};
const __findMin = (node, dim) => {
if (!node) return;
if (node.d === dim) {
return node.l ? __findMin(node.l, dim) : node;
}
const q = node.k[dim];
const l = __findMin(node.l, dim);
const r = __findMin(node.r, dim);
let min = node;
if (l && l.k[dim] < q) {
min = l;
}
if (r && r.k[dim] < min.k[dim]) {
min = r;
}
return min;
};
const __remove = (node) => {
if (!node.l && !node.r) {
if (!node.parent) {
return true;
}
const parent = node.parent;
const pdim = parent.d;
parent[node.k[pdim] < parent.k[pdim] ? "l" : "r"] = void 0;
return;
}
let next;
let nextP;
if (node.r) {
next = __findMin(node.r, node.d);
nextP = next.k;
__remove(next);
node.k = nextP;
} else {
next = __findMin(node.l, node.d);
nextP = next.k;
__remove(next);
node.r = node.l;
node.l = void 0;
node.k = nextP;
}
};
const __nearest = (q, acc, dims, maxNum, node, distFn) => {
const p = node.k;
const ndist = distSq(p, q);
if (!node.l && !node.r) {
__collect(acc, maxNum, node, ndist);
return;
}
const tdist = __nodeDist(node, dims, q, p, distFn);
let best = __bestChild(node, q);
__nearest(q, acc, dims, maxNum, best, distFn);
__collect(acc, maxNum, node, ndist);
if (tdist < acc.values[0][0]) {
best = best === node.l ? node.r : node.l;
best && __nearest(q, acc, dims, maxNum, best, distFn);
}
};
const __nearest1 = (q, acc, dims, node, distFn) => {
const p = node.k;
const ndist = distFn(p, q);
if (!node.l && !node.r) {
__collect1(acc, node, ndist);
return acc;
}
const tdist = __nodeDist(node, dims, q, p, distFn);
let best = __bestChild(node, q);
__nearest1(q, acc, dims, best, distFn);
__collect1(acc, node, ndist);
if (tdist < acc[0]) {
best = best === node.l ? node.r : node.l;
best && __nearest1(q, acc, dims, best, distFn);
}
return acc;
};
const __bestChild = (node, q) => {
const d = node.d;
return !node.r ? node.l : !node.l ? node.r : q[d] < node.k[d] ? node.l : node.r;
};
const __collect = (acc, maxNum, node, ndist) => (!acc.length || ndist < acc.peek()[0]) && (acc.length >= maxNum ? acc.pushPop([ndist, node]) : acc.push([ndist, node]));
const __collect1 = (acc, node, ndist) => ndist < acc[0] && (acc[0] = ndist, acc[1] = node);
const TMP = [];
const __nodeDist = (node, dims, q, p, distFn) => {
for (let i = dims, d = node.d; i-- > 0; ) {
TMP[i] = i === d ? q[i] : p[i];
}
return distFn(TMP, p);
};
export {
KdNode,
KdTreeMap
};