kd-tree-map.ts

import type { Fn, ICopy, IEmpty, Maybe, Pair } from "@thi.ng/api";
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 type { DistanceFn, ReadonlyVec, Vec } from "@thi.ng/vectors";
import { distSq } from "@thi.ng/vectors/distsq";
import type { IRegionQuery, ISpatialMap } from "./api.js";
import { CMP, __addResults, __into } from "./utils.js";

type MaybeKdNode<K extends ReadonlyVec, V> = Maybe<KdNode<K, V>>;

export class KdNode<K extends ReadonlyVec, V> {
	d: number;
	parent?: KdNode<K, V>;
	l?: KdNode<K, V>;
	r?: KdNode<K, V>;
	k: K;
	v: V;

	constructor(parent: MaybeKdNode<K, V>, dim: number, key: K, val: V) {
		this.parent = parent;
		this.d = dim;
		this.k = key;
		this.v = val;
	}

	get height(): number {
		return (
			1 + Math.max(this.l ? this.l.height : 0, this.r ? this.r.height : 0)
		);
	}
}

/**
 * https://en.wikipedia.org/wiki/K-d_tree
 *
 * Partially based on:
 * https://github.com/ubilabs/kd-tree-javascript
 *
 */
export class KdTreeMap<K extends ReadonlyVec, V>
	implements
		ICopy<KdTreeMap<K, V>>,
		IEmpty<KdTreeMap<K, V>>,
		IRegionQuery<K, V, number>,
		ISpatialMap<K, V>
{
	readonly dim: number;

	protected root: MaybeKdNode<K, V>;
	protected _size: number;

	constructor(
		dim: number,
		pairs?: Iterable<Pair<K, V>>,
		public readonly distanceFn: DistanceFn = distSq
	) {
		this.dim = dim;
		this._size = 0;
		this.root = pairs ? this.buildTree(ensureArray(pairs), 0) : undefined;
	}

	*[Symbol.iterator]() {
		let queue: MaybeKdNode<K, V>[] = this.root ? [this.root] : [];
		while (queue.length) {
			const n = queue.pop();
			if (n) {
				yield <Pair<K, V>>[n.k, n.v];
				queue.push(n.r, n.l);
			}
		}
	}

	*keys() {
		let queue: MaybeKdNode<K, V>[] = 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<K, V>(this.dim, undefined, this.distanceFn);
	}

	set(key: K, val: V, eps = EPS) {
		eps = Math.max(0, eps);
		eps *= eps;
		const search = (
			node: MaybeKdNode<K, V>,
			parent: MaybeKdNode<K, V>
		): MaybeKdNode<K, V> =>
			node
				? search(key[node.d] < node.k[node.d] ? node.l : node.r, node)
				: parent;
		let parent: MaybeKdNode<K, V>;
		if (this.root) {
			parent = __nearest1(
				key,
				[eps, undefined],
				this.dim,
				this.root,
				this.distanceFn
			)[1];
			if (parent) {
				parent.v = val;
				return false;
			}
			parent = search(this.root, undefined)!;
			const dim = parent.d;
			parent[key[dim] < parent.k[dim] ? "l" : "r"] = new KdNode<K, V>(
				parent,
				(dim + 1) % this.dim,
				key,
				val
			);
		} else {
			this.root = new KdNode<K, V>(undefined, 0, key, val);
		}
		this._size++;
		return true;
	}

	into(pairs: Iterable<Pair<K, V>>, eps = EPS) {
		return __into(this, pairs, eps);
	}

	remove(key: K) {
		const node = __find(key, this.root, 0);
		if (node) {
			__remove(node) && (this.root = undefined);
			this._size--;
			return true;
		}
		return false;
	}

	has(key: K, eps = EPS) {
		return (
			!!this.root &&
			!!__nearest1(
				key,
				[eps * eps, undefined],
				this.dim,
				this.root,
				this.distanceFn
			)[1]
		);
	}

	get(key: K, eps = EPS) {
		if (this.root) {
			const node = __nearest1(
				key,
				[eps * eps, undefined],
				this.dim,
				this.root,
				this.distanceFn
			)[1];
			return node ? node.v : undefined;
		}
	}

	query(
		q: K,
		maxDist: number,
		limit?: number,
		acc?: Pair<K, V>[]
	): Pair<K, V>[] {
		return this.doSelect(q, (x) => [x.k, x.v], maxDist, limit, acc);
	}

	queryKeys(q: K, maxDist: number, limit?: number, acc?: K[]): K[] {
		return this.doSelect(q, (x) => x.k, maxDist, limit, acc);
	}

	queryValues(q: K, maxDist: number, limit?: number, acc?: V[]): V[] {
		return this.doSelect(q, (x) => x.v, maxDist, limit, acc);
	}

	protected doSelect<T>(
		q: K,
		f: Fn<KdNode<K, V>, T>,
		maxDist: number,
		maxNum = 1,
		acc: T[] = []
	): T[] {
		if (!this.root) return [];
		maxDist *= maxDist;
		if (maxNum === 1) {
			const sel = __nearest1(
				q,
				[maxDist, undefined],
				this.dim,
				this.root,
				this.distanceFn
			)[1];
			sel && acc.push(f(sel));
		} else {
			const nodes = new Heap<[number, MaybeKdNode<K, V>]>(
				[[maxDist, undefined]],
				{
					compare: CMP,
				}
			);
			__nearest(q, nodes, this.dim, maxNum, this.root!, this.distanceFn);
			return __addResults(f, nodes.values, acc);
		}
		return acc;
	}

	protected buildTree(
		points: Pair<K, V>[],
		depth: number,
		parent?: KdNode<K, V>
	) {
		const n = points.length;
		if (n === 0) {
			return;
		}
		this._size++;
		let dim = depth % this.dim;
		if (n === 1) {
			return new KdNode<K, V>(parent, dim, ...points[0]);
		}
		points.sort((a, b) => a[0][dim] - b[0][dim]);
		const med = n >>> 1;
		const node = new KdNode<K, V>(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;
	}
}

/**
 * Returns node for point or `undefined` if none found.
 *
 * @param p - point
 * @param node - tree node
 * @param epsSq - squared epsilon / tolerance
 *
 * @internal
 */
const __find = <K extends ReadonlyVec, V>(
	p: K,
	node: MaybeKdNode<K, V>,
	epsSq: number
): MaybeKdNode<K, V> => {
	if (!node) return;
	return distSq(p, node.k) <= epsSq
		? node
		: __find(p, p[node.d] < node.k[node.d] ? node.l : node.r, epsSq);
};

/** @internal */
const __findMin = <K extends ReadonlyVec, V>(
	node: MaybeKdNode<K, V>,
	dim: number
): MaybeKdNode<K, V> => {
	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;
};

/**
 * Returns true if root is to be deleted.
 *
 * @param node - tree node
 *
 * @internal
 */
const __remove = <K extends ReadonlyVec, V>(node: KdNode<K, V>) => {
	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"] = undefined;
		return;
	}
	let next: MaybeKdNode<K, V>;
	let nextP: K;
	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 = undefined;
		node.k = nextP;
	}
};

/** @internal */
const __nearest = <K extends ReadonlyVec, V>(
	q: K,
	acc: Heap<[number, MaybeKdNode<K, V>]>,
	dims: number,
	maxNum: number,
	node: KdNode<K, V>,
	distFn: DistanceFn
) => {
	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);
	}
};

/**
 * Optimized version of {@link __nearest} for single closest point search.
 *
 * @param q - search point
 * @param acc - accumulator
 * @param dims - dimensions
 * @param node - tree node
 *
 * @internal
 */
const __nearest1 = <K extends ReadonlyVec, V>(
	q: K,
	acc: [number, MaybeKdNode<K, V>],
	dims: number,
	node: KdNode<K, V>,
	distFn: DistanceFn
): [number, MaybeKdNode<K, V>] => {
	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;
};

/** @internal */
const __bestChild = <K extends ReadonlyVec, V>(node: KdNode<K, V>, q: K) => {
	const d = node.d;
	return !node.r
		? node.l
		: !node.l
		? node.r
		: q[d] < node.k[d]
		? node.l
		: node.r;
};

/** @internal */
const __collect = <K extends ReadonlyVec, V>(
	acc: Heap<[number, MaybeKdNode<K, V>]>,
	maxNum: number,
	node: KdNode<K, V>,
	ndist: number
) =>
	(!acc.length || ndist < acc.peek()![0]) &&
	(acc.length >= maxNum
		? acc.pushPop([ndist, node])
		: acc.push([ndist, node]));

/** @internal */
const __collect1 = <K extends ReadonlyVec, V>(
	acc: [number, MaybeKdNode<K, V>],
	node: KdNode<K, V>,
	ndist: number
) => ndist < acc[0] && ((acc[0] = ndist), (acc[1] = node));

const TMP: Vec = [];

/** @internal */
const __nodeDist = <K extends ReadonlyVec, V>(
	node: KdNode<K, V>,
	dims: number,
	q: K,
	p: K,
	distFn: DistanceFn
) => {
	for (let i = dims, d = node.d; i-- > 0; ) {
		TMP[i] = i === d ? q[i] : p[i];
	}
	return distFn(TMP, p);
};