Separate tree nodes into two distinct types

coocos · Nov 14, 2020 · ca2e7f7 · ca2e7f7
1 parent 022264b
commit ca2e7f7
Show file tree

Hide file tree

Showing 3 changed files with 60 additions and 50 deletions.
diff --git a/src/index.ts b/src/index.ts
@@ -1,4 +1,4 @@
-import { Node, construct, nodes } from "./quadtree";
+import { construct, nodes } from "./quadtree";
 import { Vector } from "./vector";
 import { clearCanvas, initializeCanvas, drawNode, drawPoints } from "./canvas";
 

diff --git a/src/quadtree.test.ts b/src/quadtree.test.ts
@@ -15,6 +15,9 @@ describe("Quadtree", () => {
       width: 1,
       height: 1,
     });
+    if ("children" in root) {
+      fail("Flat tree root node should be a leaf");
+    }
     expect(root.points).not.toEqual(undefined);
     expect(root.points).toEqual(points);
   });
@@ -32,8 +35,10 @@ describe("Quadtree", () => {
       width: 1,
       height: 1,
     });
-    expect(root.points).toEqual(undefined);
-    expect(root.children?.topLeft).toEqual({
+    if ("points" in root) {
+      fail("Two-level tree root should not be a leaf");
+    }
+    expect(root.children.topLeft).toEqual({
       box: {
         x: 0,
         y: 0,
@@ -42,7 +47,7 @@ describe("Quadtree", () => {
       },
       points: [points[0]],
     });
-    expect(root.children?.topRight).toEqual({
+    expect(root.children.topRight).toEqual({
       box: {
         x: 0.5,
         y: 0.0,
@@ -51,7 +56,7 @@ describe("Quadtree", () => {
       },
       points: [points[1]],
     });
-    expect(root.children?.bottomLeft).toEqual({
+    expect(root.children.bottomLeft).toEqual({
       box: {
         x: 0.0,
         y: 0.5,
@@ -60,7 +65,7 @@ describe("Quadtree", () => {
       },
       points: [points[2]],
     });
-    expect(root.children?.bottomRight).toEqual({
+    expect(root.children.bottomRight).toEqual({
       box: {
         x: 0.5,
         y: 0.5,

diff --git a/src/quadtree.ts b/src/quadtree.ts
@@ -1,34 +1,72 @@
 import { Vector } from "./vector";
 
+const BUCKET_SIZE = 4;
+
 type BoundingBox = {
   x: number;
   y: number;
   width: number;
   height: number;
 };
 
-export type Node = {
-  children?: {
+type Leaf = {
+  box: BoundingBox;
+  points: Vector[];
+};
+
+type Inner = {
+  box: BoundingBox;
+  children: {
     topLeft: Node;
     topRight: Node;
     bottomLeft: Node;
     bottomRight: Node;
   };
-  box: BoundingBox;
-  points?: Vector[];
 };
 
+type Node = Leaf | Inner;
+
 export function construct(points: Vector[], box: BoundingBox): Node {
-  const root = {
+  let root: Node = {
     box,
     points: [],
   };
   for (let point of points) {
-    insert(root, point);
+    root = insert(root, point);
   }
   return root;
 }
-const BUCKET_SIZE = 4;
+
+function insert(node: Node, point: Vector): Node {
+  if (isLeaf(node)) {
+    if (node.points.length < BUCKET_SIZE) {
+      node.points.push(point);
+      return node;
+    }
+    const inner: Inner = {
+      box: node.box,
+      children: split(node),
+    };
+    for (let nodePoint of [...node.points, point]) {
+      for (let [name, child] of Object.entries(inner.children)) {
+        if (intersects(child, nodePoint)) {
+          inner.children[name as keyof Inner["children"]] = insert(
+            child,
+            nodePoint
+          );
+        }
+      }
+    }
+    return inner;
+  } else {
+    for (let [name, child] of Object.entries(node.children)) {
+      if (intersects(child, point)) {
+        node.children[name as keyof Inner["children"]] = insert(child, point);
+      }
+    }
+    return node;
+  }
+}
 
 function intersects(node: Node, point: Vector) {
   return (
@@ -39,7 +77,7 @@ function intersects(node: Node, point: Vector) {
   );
 }
 
-export function split(node: Node) {
+function split(node: Leaf) {
   const childWidth = node.box.width / 2;
   const childHeight = node.box.height / 2;
   const dimensions = {
@@ -82,14 +120,13 @@ export function split(node: Node) {
   };
 }
 
-// FIXME: You should really use *[Symbol.Iterator somehow]?
 export function nodes(node: Node) {
   const queue = [node];
   const nodes = [];
   while (queue.length > 0) {
     const node = queue.shift()!;
     nodes.push(node);
-    if (node.children) {
+    if (!isLeaf(node)) {
       for (let child of Object.values(node.children)) {
         queue.push(child);
       }
@@ -98,38 +135,6 @@ export function nodes(node: Node) {
   return nodes;
 }
 
-function isLeaf(node: Node) {
-  return node.children === undefined && node.points !== undefined;
-}
-
-export function insert(node: Node, point: Vector) {
-  if (isLeaf(node)) {
-    if (node.points !== undefined) {
-      if (node.points.length < BUCKET_SIZE) {
-        node.points.push(point);
-      } else {
-        node.children = split(node);
-        for (let nodePoint of [...node.points, point]) {
-          for (let child of Object.values(node.children)) {
-            if (intersects(child, nodePoint)) {
-              insert(child, nodePoint);
-            }
-          }
-        }
-        node.points = undefined;
-      }
-    } else {
-      console.error("Leaf should always have points");
-    }
-  } else {
-    if (node.children) {
-      for (let child of Object.values(node.children)) {
-        if (intersects(child, point)) {
-          insert(child, point);
-        }
-      }
-    } else {
-      console.error("Non-leaf should always have children");
-    }
-  }
+function isLeaf(node: Node): node is Leaf {
+  return "points" in node;
 }