Add function for finding points within a given radius

src/index.ts

@@ -1,4 +1,4 @@
-import { insert, construct } from "./quadtree";
+import { insert, construct, pointsWithinArea } from "./quadtree";
 import { initializeCanvas } from "./canvas";
 
 const canvas = initializeCanvas();

src/quadtree.test.ts

@@ -1,4 +1,4 @@
-import { construct } from "./quadtree";
+import { construct, pointsWithinArea } from "./quadtree";
 
 describe("Quadtree", () => {
   test("constructing a flat tree", () => {
@@ -101,4 +101,41 @@ describe("Quadtree", () => {
       points: [points[3], points[4]],
     });
   });
+  test("finding points within an area", () => {
+    const points = [
+      {
+        x: 0,
+        y: 0,
+      },
+      {
+        x: 1,
+        y: 0,
+      },
+      {
+        x: 0,
+        y: 1,
+      },
+      {
+        x: 1,
+        y: 1,
+      },
+      {
+        x: 9,
+        y: 9,
+      },
+      {
+        x: 8,
+        y: 8,
+      },
+    ];
+    const root = construct(points, {
+      x: 0,
+      y: 0,
+      width: 10,
+      height: 10,
+    });
+    expect(pointsWithinArea(root, { x: 1, y: 1, radius: 2 })).toEqual(
+      points.slice(0, 4)
+    );
+  });
 });

src/quadtree.ts

@@ -5,6 +5,10 @@ type Point = {
   y: number;
 };
 
+type Area = Point & {
+  radius: number;
+};
+
 type BoundingBox = Point & {
   width: number;
   height: number;
@@ -38,6 +42,19 @@ export function construct(points: Point[], box: BoundingBox): Node {
   return root;
 }
 
+export function pointsWithinArea(node: Node, area: Area) {
+  if (isLeaf(node)) {
+    return node.points.filter((point) => pointWithinArea(point, area));
+  }
+  let points: Point[] = [];
+  for (let child of Object.values(node.children)) {
+    if (boxWithinArea(child.box, area)) {
+      points = [...points, ...pointsWithinArea(child, area)];
+    }
+  }
+  return points;
+}
+
 export function insert(node: Node, point: Point): Node {
   if (isLeaf(node)) {
     if (node.points.length < BUCKET_SIZE) {
@@ -69,6 +86,25 @@ export function insert(node: Node, point: Point): Node {
   }
 }
 
+export function nodes(node: Node) {
+  const queue = [node];
+  const nodes = [];
+  while (queue.length > 0) {
+    const node = queue.shift()!;
+    nodes.push(node);
+    if (!isLeaf(node)) {
+      for (let child of Object.values(node.children)) {
+        queue.push(child);
+      }
+    }
+  }
+  return nodes;
+}
+
+export function isLeaf(node: Node): node is Leaf {
+  return "points" in node;
+}
+
 function intersects(node: Node, point: Point) {
   return (
     point.x <= node.box.x + node.box.width &&
@@ -121,21 +157,14 @@ function split(node: Leaf) {
   };
 }
 
-export function nodes(node: Node) {
-  const queue = [node];
-  const nodes = [];
-  while (queue.length > 0) {
-    const node = queue.shift()!;
-    nodes.push(node);
-    if (!isLeaf(node)) {
-      for (let child of Object.values(node.children)) {
-        queue.push(child);
-      }
-    }
-  }
-  return nodes;
+function boxWithinArea(box: BoundingBox, area: Area) {
+  const closestX = Math.max(box.x, Math.min(area.x, box.x + box.width));
+  const closestY = Math.max(box.y, Math.min(area.y, box.y + box.height));
+  const distanceSquared = (area.x - closestX) ** 2 + (area.y - closestY) ** 2;
+  return distanceSquared < area.radius ** 2;
 }
 
-export function isLeaf(node: Node): node is Leaf {
-  return "points" in node;
+function pointWithinArea(point: Point, area: Area) {
+  const distanceSquared = (point.x - area.x) ** 2 + (point.y - area.y) ** 2;
+  return distanceSquared < area.radius ** 2;
 }