Memorization optimziations

src/kd-tree.ts

@@ -15,6 +15,7 @@ interface KDNode<T> {
 	axis: number;
 	left: KDNode<T> | null;
 	right: KDNode<T> | null;
+	deleted?: boolean;
 }
 
 // ─── Distance helpers ─────────────────────────────────────────────────────────
@@ -95,6 +96,7 @@ class BoundedMaxHeap<T> {
  *
  * Supports:
  *  - Insertion of labeled points
+ *  - Lazy (tombstone) removal, physically purged on rebalance()
  *  - k-nearest-neighbor (KNN) search
  *  - Radius search (all points within a given distance)
  *  - Euclidean and cosine distance metrics
@@ -103,6 +105,7 @@ class BoundedMaxHeap<T> {
 export class KDTree<T = unknown> {
 	private root: KDNode<T> | null = null;
 	private _size = 0;
+	private _tombstones = 0;
 	private readonly distanceFn: (a: number[], b: number[]) => number;
 
 	readonly dims: number;
@@ -129,9 +132,15 @@ export class KDTree<T = unknown> {
 		}
 	}
 
-	/** Total number of points stored in the tree. */
+	/** Total number of live points stored in the tree (excludes tombstoned). */
 	get size(): number { return this._size; }
 
+	/** Fraction of physical nodes that are tombstoned (pending removal on next rebalance). */
+	get tombstoneRatio(): number {
+		const total = this._size + this._tombstones;
+		return total ? this._tombstones / total : 0;
+	}
+
 	// ── Insertion ──────────────────────────────────────────────────────────────
 
 	/**
@@ -144,10 +153,36 @@ export class KDTree<T = unknown> {
 		this._size++;
 	}
 
+	// ── Removal ────────────────────────────────────────────────────────────────
+
+	/**
+	 * Lazily remove all live points whose payload matches `predicate`.
+	 * O(n) traversal, but avoids a full tree rebuild. Call `rebalance()`
+	 * periodically (e.g. once tombstoneRatio crosses ~0.25) to reclaim space
+	 * and restore optimal query depth.
+	 * @returns number of points removed
+	 */
+	remove(predicate: (payload: T) => boolean): number {
+		let removed = 0;
+		const visit = (node: KDNode<T> | null): void => {
+			if (!node) return;
+			if (!node.deleted && predicate(node.point.payload)) {
+				node.deleted = true;
+				removed++;
+			}
+			visit(node.left);
+			visit(node.right);
+		};
+		visit(this.root);
+		this._size -= removed;
+		this._tombstones += removed;
+		return removed;
+	}
+
 	// ── KNN search ─────────────────────────────────────────────────────────────
 
 	/**
-	 * Find the k nearest neighbors to `query`.
+	 * Find the k nearest live neighbors to `query`.
 	 * Returns results sorted by distance ascending.
 	 */
 	knn(query: number[], k: number): KNNResult<T>[] {
@@ -171,7 +206,7 @@ export class KDTree<T = unknown> {
 	// ── Radius search ──────────────────────────────────────────────────────────
 
 	/**
-	 * Return all points whose distance to `query` is ≤ `radius`,
+	 * Return all live points whose distance to `query` is ≤ `radius`,
 	 * sorted by distance ascending.
 	 */
 	radiusSearch(query: number[], radius: number): KNNResult<T>[] {
@@ -186,7 +221,7 @@ export class KDTree<T = unknown> {
 
 	// ── Conversion ─────────────────────────────────────────────────────────────
 
-	/** Collect all points in the tree (order not guaranteed). */
+	/** Collect all live points in the tree (order not guaranteed). */
 	toArray(): KDPoint<T>[] {
 		const out: KDPoint<T>[] = [];
 		this.collect(this.root, out);
@@ -194,12 +229,14 @@ export class KDTree<T = unknown> {
 	}
 
 	/**
-	 * Rebuild the tree from its current points as a balanced tree.
-	 * Useful after many individual insertions to restore O(log n) query time.
+	 * Rebuild the tree from its current live points as a balanced tree.
+	 * Physically purges tombstones and restores O(log n) query time.
 	 */
 	rebalance(): void {
 		const points = this.toArray();
 		this.root = points.length ? this.buildBalanced(points, 0) : null;
+		this._size = points.length;
+		this._tombstones = 0;
 	}
 
 	// ── Private: build ─────────────────────────────────────────────────────────
@@ -251,8 +288,10 @@ export class KDTree<T = unknown> {
 	): void {
 		if (node === null) return;
 
-		const dist = this.distanceFn(query, node.point.vector);
-		heap.push({ point: node.point, distance: dist });
+		if (!node.deleted) {
+			const dist = this.distanceFn(query, node.point.vector);
+			heap.push({ point: node.point, distance: dist });
+		}
 
 		const axis = node.axis;
 		const diff = query[axis] - node.point.vector[axis];
@@ -285,9 +324,11 @@ export class KDTree<T = unknown> {
 	): void {
 		if (node === null) return;
 
-		const dist = this.distanceFn(query, node.point.vector);
-		if (dist <= radius) {
-			results.push({ point: node.point, distance: dist });
+		if (!node.deleted) {
+			const dist = this.distanceFn(query, node.point.vector);
+			if (dist <= radius) {
+				results.push({ point: node.point, distance: dist });
+			}
 		}
 
 		const axis = node.axis;
@@ -310,7 +351,7 @@ export class KDTree<T = unknown> {
 
 	private collect(node: KDNode<T> | null, out: KDPoint<T>[]): void {
 		if (node === null) return;
-		out.push(node.point);
+		if (!node.deleted) out.push(node.point);
 		this.collect(node.left, out);
 		this.collect(node.right, out);
 	}