336 lines
10 KiB
TypeScript
336 lines
10 KiB
TypeScript
export type DistanceMetric = "euclidean" | "cosine";
|
|
|
|
export interface KDPoint<T = unknown> {
|
|
vector: number[];
|
|
payload: T;
|
|
}
|
|
|
|
export interface KNNResult<T = unknown> {
|
|
point: KDPoint<T>;
|
|
distance: number;
|
|
}
|
|
|
|
interface KDNode<T> {
|
|
point: KDPoint<T>;
|
|
axis: number;
|
|
left: KDNode<T> | null;
|
|
right: KDNode<T> | null;
|
|
}
|
|
|
|
// ─── Distance helpers ─────────────────────────────────────────────────────────
|
|
|
|
function euclidean(a: number[], b: number[]): number {
|
|
let sum = 0;
|
|
for (let i = 0; i < a.length; i++) {
|
|
const d = a[i] - b[i];
|
|
sum += d * d;
|
|
}
|
|
return Math.sqrt(sum);
|
|
}
|
|
|
|
function cosine(a: number[], b: number[]): number {
|
|
let dot = 0, normA = 0, normB = 0;
|
|
for (let i = 0; i < a.length; i++) {
|
|
dot += a[i] * b[i];
|
|
normA += a[i] * a[i];
|
|
normB += b[i] * b[i];
|
|
}
|
|
const denom = Math.sqrt(normA) * Math.sqrt(normB);
|
|
return denom === 0 ? 1 : 1 - dot / denom; // distance = 1 - similarity
|
|
}
|
|
|
|
/**
|
|
* Keeps the k closest candidates in memory, evicts the furthest when full
|
|
*/
|
|
class BoundedMaxHeap<T> {
|
|
private heap: KNNResult<T>[] = [];
|
|
|
|
constructor(private readonly k: number) {}
|
|
|
|
get size(): number { return this.heap.length; }
|
|
|
|
get worstDistance(): number {
|
|
return this.heap.length < this.k ? Infinity : this.heap[0].distance;
|
|
}
|
|
|
|
push(item: KNNResult<T>): void {
|
|
if (this.heap.length < this.k) {
|
|
this.heap.push(item);
|
|
this.bubbleUp(this.heap.length - 1);
|
|
} else if (item.distance < this.heap[0].distance) {
|
|
this.heap[0] = item;
|
|
this.sinkDown(0);
|
|
}
|
|
}
|
|
|
|
toSortedArray(): KNNResult<T>[] {
|
|
return [...this.heap].sort((a, b) => a.distance - b.distance);
|
|
}
|
|
|
|
private bubbleUp(i: number): void {
|
|
while (i > 0) {
|
|
const parent = (i - 1) >> 1;
|
|
if (this.heap[parent].distance >= this.heap[i].distance) break;
|
|
[this.heap[parent], this.heap[i]] = [this.heap[i], this.heap[parent]];
|
|
i = parent;
|
|
}
|
|
}
|
|
|
|
private sinkDown(i: number): void {
|
|
const n = this.heap.length;
|
|
while (true) {
|
|
let largest = i;
|
|
const l = 2 * i + 1, r = 2 * i + 2;
|
|
if (l < n && this.heap[l].distance > this.heap[largest].distance) largest = l;
|
|
if (r < n && this.heap[r].distance > this.heap[largest].distance) largest = r;
|
|
if (largest === i) break;
|
|
[this.heap[largest], this.heap[i]] = [this.heap[i], this.heap[largest]];
|
|
i = largest;
|
|
}
|
|
}
|
|
}
|
|
|
|
/**
|
|
* K-D Tree for efficient nearest-neighbor search over high-dimensional vectors / embeddings.
|
|
*
|
|
* Supports:
|
|
* - Insertion of labeled points
|
|
* - k-nearest-neighbor (KNN) search
|
|
* - Radius search (all points within a given distance)
|
|
* - Euclidean and cosine distance metrics
|
|
* - Bulk construction (balanced tree) for best query performance
|
|
*/
|
|
export class KDTree<T = unknown> {
|
|
private root: KDNode<T> | null = null;
|
|
private _size = 0;
|
|
private readonly distanceFn: (a: number[], b: number[]) => number;
|
|
|
|
readonly dims: number;
|
|
|
|
/**
|
|
* @param dims Dimensionality of all vectors (must be consistent).
|
|
* @param metric Distance metric to use. Default: "euclidean".
|
|
* @param points Optional initial set of points. Builds a balanced tree
|
|
* in O(n log² n) — prefer this over inserting one-by-one
|
|
* when you have a large corpus.
|
|
*/
|
|
constructor(
|
|
dims: number,
|
|
metric: DistanceMetric = "euclidean",
|
|
points?: KDPoint<T>[]
|
|
) {
|
|
this.dims = dims;
|
|
this.distanceFn = metric === "cosine" ? cosine : euclidean;
|
|
|
|
if (points && points.length > 0) {
|
|
this.validateAll(points);
|
|
this.root = this.buildBalanced([...points], 0);
|
|
this._size = points.length;
|
|
}
|
|
}
|
|
|
|
/** Total number of points stored in the tree. */
|
|
get size(): number { return this._size; }
|
|
|
|
// ── Insertion ──────────────────────────────────────────────────────────────
|
|
|
|
/**
|
|
* Insert a single point. O(log n) average, O(n) worst case on skewed data.
|
|
* For bulk loading prefer passing points to the constructor.
|
|
*/
|
|
insert(point: KDPoint<T>): void {
|
|
this.validate(point);
|
|
this.root = this.insertNode(this.root, point, 0);
|
|
this._size++;
|
|
}
|
|
|
|
// ── KNN search ─────────────────────────────────────────────────────────────
|
|
|
|
/**
|
|
* Find the k nearest neighbors to `query`.
|
|
* Returns results sorted by distance ascending.
|
|
*/
|
|
knn(query: number[], k: number): KNNResult<T>[] {
|
|
if (k <= 0) throw new RangeError("k must be a positive integer");
|
|
this.validateVector(query);
|
|
|
|
const heap = new BoundedMaxHeap<T>(k);
|
|
this.searchKNN(this.root, query, k, heap, 0);
|
|
return heap.toSortedArray();
|
|
}
|
|
|
|
/**
|
|
* Nearest single neighbor. Convenience wrapper around knn(query, 1).
|
|
* Returns null if the tree is empty.
|
|
*/
|
|
nearest(query: number[]): KNNResult<T> | null {
|
|
const results = this.knn(query, 1);
|
|
return results[0] ?? null;
|
|
}
|
|
|
|
// ── Radius search ──────────────────────────────────────────────────────────
|
|
|
|
/**
|
|
* Return all points whose distance to `query` is ≤ `radius`,
|
|
* sorted by distance ascending.
|
|
*/
|
|
radiusSearch(query: number[], radius: number): KNNResult<T>[] {
|
|
if (radius < 0) throw new RangeError("radius must be non-negative");
|
|
this.validateVector(query);
|
|
|
|
const results: KNNResult<T>[] = [];
|
|
this.searchRadius(this.root, query, radius, results, 0);
|
|
results.sort((a, b) => a.distance - b.distance);
|
|
return results;
|
|
}
|
|
|
|
// ── Conversion ─────────────────────────────────────────────────────────────
|
|
|
|
/** Collect all points in the tree (order not guaranteed). */
|
|
toArray(): KDPoint<T>[] {
|
|
const out: KDPoint<T>[] = [];
|
|
this.collect(this.root, out);
|
|
return out;
|
|
}
|
|
|
|
/**
|
|
* Rebuild the tree from its current points as a balanced tree.
|
|
* Useful after many individual insertions to restore O(log n) query time.
|
|
*/
|
|
rebalance(): void {
|
|
const points = this.toArray();
|
|
this.root = points.length ? this.buildBalanced(points, 0) : null;
|
|
}
|
|
|
|
// ── Private: build ─────────────────────────────────────────────────────────
|
|
|
|
private buildBalanced(points: KDPoint<T>[], depth: number): KDNode<T> {
|
|
const axis = depth % this.dims;
|
|
points.sort((a, b) => a.vector[axis] - b.vector[axis]);
|
|
|
|
const mid = Math.floor(points.length / 2);
|
|
return {
|
|
point: points[mid],
|
|
axis,
|
|
left: points.slice(0, mid).length
|
|
? this.buildBalanced(points.slice(0, mid), depth + 1)
|
|
: null,
|
|
right: points.slice(mid + 1).length
|
|
? this.buildBalanced(points.slice(mid + 1), depth + 1)
|
|
: null,
|
|
};
|
|
}
|
|
|
|
// ── Private: insert ────────────────────────────────────────────────────────
|
|
|
|
private insertNode(
|
|
node: KDNode<T> | null,
|
|
point: KDPoint<T>,
|
|
depth: number
|
|
): KDNode<T> {
|
|
if (node === null) {
|
|
return { point, axis: depth % this.dims, left: null, right: null };
|
|
}
|
|
const axis = depth % this.dims;
|
|
if (point.vector[axis] < node.point.vector[axis]) {
|
|
node.left = this.insertNode(node.left, point, depth + 1);
|
|
} else {
|
|
node.right = this.insertNode(node.right, point, depth + 1);
|
|
}
|
|
return node;
|
|
}
|
|
|
|
// ── Private: KNN traversal ─────────────────────────────────────────────────
|
|
|
|
private searchKNN(
|
|
node: KDNode<T> | null,
|
|
query: number[],
|
|
k: number,
|
|
heap: BoundedMaxHeap<T>,
|
|
depth: number
|
|
): void {
|
|
if (node === null) return;
|
|
|
|
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];
|
|
const [near, far] = diff <= 0
|
|
? [node.left, node.right]
|
|
: [node.right, node.left];
|
|
|
|
this.searchKNN(near, query, k, heap, depth + 1);
|
|
|
|
// Only explore the far side if it could contain a closer point.
|
|
// For cosine distance we can't prune by axis gap alone, so always explore.
|
|
const shouldExplore =
|
|
this.distanceFn === cosine
|
|
? true
|
|
: Math.abs(diff) < heap.worstDistance;
|
|
|
|
if (shouldExplore) {
|
|
this.searchKNN(far, query, k, heap, depth + 1);
|
|
}
|
|
}
|
|
|
|
// ── Private: radius traversal ──────────────────────────────────────────────
|
|
|
|
private searchRadius(
|
|
node: KDNode<T> | null,
|
|
query: number[],
|
|
radius: number,
|
|
results: KNNResult<T>[],
|
|
depth: number
|
|
): void {
|
|
if (node === null) return;
|
|
|
|
const dist = this.distanceFn(query, node.point.vector);
|
|
if (dist <= radius) {
|
|
results.push({ point: node.point, distance: dist });
|
|
}
|
|
|
|
const axis = node.axis;
|
|
const diff = query[axis] - node.point.vector[axis];
|
|
const [near, far] = diff <= 0
|
|
? [node.left, node.right]
|
|
: [node.right, node.left];
|
|
|
|
this.searchRadius(near, query, radius, results, depth + 1);
|
|
|
|
const shouldExplore =
|
|
this.distanceFn === cosine ? true : Math.abs(diff) <= radius;
|
|
|
|
if (shouldExplore) {
|
|
this.searchRadius(far, query, radius, results, depth + 1);
|
|
}
|
|
}
|
|
|
|
// ── Private: collect ───────────────────────────────────────────────────────
|
|
|
|
private collect(node: KDNode<T> | null, out: KDPoint<T>[]): void {
|
|
if (node === null) return;
|
|
out.push(node.point);
|
|
this.collect(node.left, out);
|
|
this.collect(node.right, out);
|
|
}
|
|
|
|
// ── Private: validation ────────────────────────────────────────────────────
|
|
|
|
private validateVector(v: number[]): void {
|
|
if (v.length !== this.dims) {
|
|
throw new TypeError(
|
|
`Vector length ${v.length} does not match tree dimensionality ${this.dims}`
|
|
);
|
|
}
|
|
}
|
|
|
|
private validate(point: KDPoint<T>): void {
|
|
this.validateVector(point.vector);
|
|
}
|
|
|
|
private validateAll(points: KDPoint<T>[]): void {
|
|
for (const p of points) this.validate(p);
|
|
}
|
|
}
|