add QTNode.remove()

2 files changed, 82 insertions, 1 deletions

diff --git a/src/Geometry/Geometry.js b/src/Geometry/Geometry.js
index 43f277f..6c296ad 100644
--- a/src/Geometry/Geometry.js
+++ b/src/Geometry/Geometry.js
@@ -122,6 +122,41 @@ class QTNode {
 		return false;
 	}
 
+	/* remove a node from the tree */
+	remove(point) {
+		if (this.aabb.contains(point)) {
+			switch(this.type) {
+				case QTNode.Empty:
+					break;
+
+				case QTNode.Leaf:
+					if (this.point.x === point.x && this.point.y === point.y) {
+						this.point = undefined;
+						this.type = QTNode.Empty;
+					}
+					break;
+
+				case QTNode.Branch:
+					for (let subnode of this.subnode) subnode.remove(point);
+					let numBranches = 0;
+					let numLeaves = 0;
+					for (let subnode of this.subnode) {
+						if (subnode.type === QTNode.Leaf) numLeaves += 1;
+						if (subnode.type === QTNode.Branch) numBranches += 1;
+					}
+					if (numBranches == 0 && numLeaves == 1) {
+						/* contract to leaf node */
+						this.type = QTNode.Empty; // start as Empty so insertion works correctly
+						for (let subnode of this.subnode) {
+							if (subnode.type == QTNode.Leaf) this.insert(subnode.point);
+						}
+						this.subnode = undefined;
+					}
+					break;
+			}
+		}
+	}
+
 	/* subdivide the node into four child nodes
 	 *
 	 * under normal circumstances, the user should not call this function --
diff --git a/src/Geometry/Geometry.test.js b/src/Geometry/Geometry.test.js
index 0b81438..1ec50f7 100644
--- a/src/Geometry/Geometry.test.js
+++ b/src/Geometry/Geometry.test.js
@@ -176,7 +176,53 @@ test('QTNode correctly returns points in region', () => {
 	};
 	bf_points.sort(compare);
 	qt_points.sort(compare);
-	assert.deepEqual(bf_points, qt_points);});
+	assert.deepEqual(bf_points, qt_points);
+});
+
+
+test('QTNode correctly removes immediate child point', () => {
+	const node = new QTNode(0, 0, 1, 1);
+
+	const pt = randomPoint();
+	node.insert(pt);
+	assert.deepEqual(node.type, QTNode.Leaf);
+	
+	const pt2 = randomPoint();
+	node.remove(pt2);
+	assert.deepEqual(node.type, QTNode.Leaf);
+	assert.deepEqual(node.point, pt);
+
+	node.remove(pt);
+	assert.deepEqual(node.type, QTNode.Empty);
+	assert.equal(node.point, undefined);
+});
+
+
+test('QTNode correctly removes point from branch', () => {
+	for (let i=0; i<100; i++) {
+		const node = new QTNode(0, 0, 1, 1);
+
+		const pt1 = randomPoint();
+		const pt2 = randomPoint();
+		const pt3 = randomPoint();
+
+		node.insert(pt1); node.insert(pt2); node.insert(pt3);
+
+		assert.deepEqual(node.type, QTNode.Branch);
+
+		node.remove(pt1);
+		// should still be a branch because there are two child points yet
+		assert.deepEqual(node.type, QTNode.Branch);
+
+		node.remove(pt2);
+		// just one now, should be a leaf
+		assert.deepEqual(node.type, QTNode.Leaf);
+
+		node.remove(pt3);
+		// empty
+		assert.deepEqual(node.type, QTNode.Empty);
+	}
+});
 
 
 test('QuadTree correctly finds closest point for 10,000 random points', () => {