implement quadtree

2 files changed, 173 insertions, 3 deletions

diff --git a/modules/Geometry.js b/modules/Geometry.js
index f32dc0e..48e682a 100644
--- a/modules/Geometry.js
+++ b/modules/Geometry.js
@@ -91,6 +91,64 @@ class QTNode {
 			new QTNode(x+w, y+h, w, h),
 		];
 	}
+
+	getPointsInRegion(aabb) {
+		if (this.type == QTNode.Empty) return [];
+		if (!this.aabb.intersects(aabb)) return [];
+		if (this.type == QTNode.Leaf) return [this.point];
+		
+		return [
+			...this.subnode[0].getPointsInRegion(aabb),
+			...this.subnode[1].getPointsInRegion(aabb),
+			...this.subnode[2].getPointsInRegion(aabb),
+			...this.subnode[3].getPointsInRegion(aabb),
+		];
+	}
+}
+
+
+class QuadTree {
+	constructor(width, height) {
+		this.root = new QTNode(0, 0, width, height);
+	}
+
+	insert(point) {
+		return this.root.insert(point);
+	}
+
+	closest(point) {
+		const getLeaf = (point, node) => {
+			if (!node.aabb.contains(point)) return null;
+			if (node.type == QTNode.Empty ||
+			    node.type == QTNode.Leaf) return node;
+			
+			for (let subnode of node.subnode) {
+				const leaf = getLeaf(point, subnode);
+				if (leaf !== null) return leaf;
+			}
+
+			return null;
+		}
+
+		const leaf = getLeaf(point, this.root);
+		if (leaf === null) return leaf;
+
+		const dim = 2*Math.max(leaf.aabb.width, leaf.aabb.height);
+		const region = new AABB(point.x - dim, point.y - dim, 2*dim, 2*dim);
+		const points = this.root.getPointsInRegion(region);
+
+		let closest = null;
+		let min_dist = 100 * this.root.aabb.width * this.root.aabb.height;
+		const dist = (a, b) => (a.x - b.x)**2 + (a.y - b.y)**2;
+		for (let pt of points) {
+			const d = dist(point, pt);
+			if (d < min_dist) {
+				closest = pt;
+				min_dist = d;
+			}
+		}
+		return closest;
+	}
 }
 
-export { AABB, QTNode };
+export { AABB, QTNode, QuadTree };
diff --git a/modules/Geometry.test.js b/modules/Geometry.test.js
index b44f25d..bfac85f 100644
--- a/modules/Geometry.test.js
+++ b/modules/Geometry.test.js
@@ -1,6 +1,6 @@
 import { test, assert } from './test-assert.js';
 
-import { AABB, QTNode } from './Geometry.js';
+import { AABB, QTNode, QuadTree } from './Geometry.js';
 
 
 test('AABB correctly contains/excludes points', () => {
@@ -141,5 +141,117 @@ test('QTNode correctly inserts multiple points', () => {
 	assert.equal(node.subnode[2].type.toString(), 'QTNode.Empty');
 
 	assert.equal(node.subnode[3].type.toString(), 'QTNode.Leaf');
-	assert.deepEqual(node.subnode[3].point, p0);
+assert.deepEqual(node.subnode[3].point, p0);
 });
+
+
+const randomPoint = () => ({ x: Math.random(), y: Math.random() });
+
+
+test('QTNode correctly returns points in region', () => {
+	const node = new QTNode(0, 0, 1, 1);
+	let points = [];
+	for (let i=0; i<100; i++) {
+		const pt = randomPoint();
+		node.insert(pt);
+		points.push(pt);
+	}
+
+	const pointsInRegion = aabb => {
+		let pts = [];
+		for (let pt of points) {
+			if (aabb.contains(pt)) pts.push(pt);
+		}
+		return pts;
+	};
+
+	const region = new AABB(0.25, 0.25, 0.5, 0.5);
+	let bf_points = pointsInRegion(region);
+	let qt_points = node.getPointsInRegion(region);
+	const compare = (a, b) => {
+		if (a.x == b.x)
+			return b.y - a.y;
+		return b.x - a.x;
+	};
+	bf_points.sort(compare);
+	qt_points.sort(compare);
+	assert.deepEqual(bf_points, qt_points);});
+
+
+test('QuadTree correctly finds closest point for 10,000 random points', () => {
+	const tree = new QuadTree(1, 1);
+
+	let points = [];
+	for (let i=0; i<10000; i++) {
+		const pt = randomPoint();
+		tree.insert(pt);
+		points.push(pt);
+	}
+	
+	const distance = (a, b) => Math.sqrt((a.x - b.x)**2 + (a.y - b.y)**2);
+	const bf_closest = a => {
+		let closest = null;
+		let min_dist = 1000;
+		for (let b of points) {
+			let d = distance(a, b);
+			if (d < min_dist) {
+				min_dist = d;
+				closest = b;
+			}
+		}
+		return closest;
+	}
+
+	for (let i=0; i<10000; i++) {
+		const pt = randomPoint();
+		assert.deepEqual(bf_closest(pt), tree.closest(pt));
+	}
+});
+
+
+function benchmark(desc, f, iterations) {
+	process.stdout.write(`${desc}: `);
+	const start = process.uptime();
+	for (let i=0; i<iterations; i++) f();
+	const elapsed = process.uptime() - start;
+	console.log(`${elapsed} seconds`);
+}
+
+const closest_benchmark = () => {
+	const n_points = 1e5;
+	const iterations = 1e4;
+	console.log(`> benchmark get closest of ${n_points} random points ${iterations} times`);
+
+	const tree = new QuadTree(1, 1);
+	let points = [];
+	for (let i=0; i<n_points; i++) {
+		const pt = randomPoint();
+		tree.insert(pt);
+		points.push(pt);
+	}
+
+	const distance = (a, b) => Math.sqrt((a.x - b.x)**2 + (a.y - b.y)**2);
+	const bf_closest = a => {
+		let closest = null;
+		let min_dist = 1000;
+		for (let b of points) {
+			let d = distance(a, b);
+			if (d < min_dist) {
+				min_dist = d;
+				closest = b;
+			}
+		}
+		return closest;
+	}
+
+	benchmark('\tbrute force', () => {
+		const pt = randomPoint();
+		bf_closest(pt);
+	}, iterations);
+
+	benchmark('\tquadtree', () => {
+		const pt = randomPoint();
+		tree.closest(pt);
+	}, iterations);
+}
+closest_benchmark();