feat: add callback functionality for binary search tree

data_structures/_test_utils.ts

@@ -8,4 +8,5 @@ export class MyMath {
 export interface Container {
   id: number;
   values: number[];
+  st_size: number;
 }

data_structures/binary_search_tree.ts

@@ -93,6 +93,7 @@ type Direction = "left" | "right";
 export class BinarySearchTree<T> implements Iterable<T> {
   #root: BinarySearchNode<T> | null = null;
   #size = 0;
+  #callback: ((node: BinarySearchNode<T>) => void) | null = null;
   #compare: (a: T, b: T) => number;
 
   /**
@@ -104,12 +105,18 @@ export class BinarySearchTree<T> implements Iterable<T> {
    * @param compare A custom comparison function to sort the values in the tree.
    * By default, the values are sorted in ascending order.
    */
-  constructor(compare: (a: T, b: T) => number = ascend) {
+  constructor(compare: (a: T, b: T) => number = ascend, callback?: (node: BinarySearchNode<T>) => void) {
     if (typeof compare !== "function") {
       throw new TypeError(
         "Cannot construct a BinarySearchTree: the 'compare' parameter is not a function, did you mean to call BinarySearchTree.from?",
       );
     }
+    if (callback && typeof callback !== "function") {
+      throw new TypeError(
+        "Cannot construct a BinarySearchTree: the 'callback' parameter is not a function.",
+      );
+    }
+    this.#callback = callback || null;
     this.#compare = compare;
   }
 
@@ -353,6 +360,10 @@ export class BinarySearchTree<T> implements Iterable<T> {
     }
     replacement[direction] = node;
     node.parent = replacement;
+    if (this.#callback) {
+      this.#callback(node);
+      this.#callback(replacement);
+    }
   }
 
   #insertNode(
@@ -374,6 +385,14 @@ export class BinarySearchTree<T> implements Iterable<T> {
         } else {
           node[direction] = new Node(node, value);
           this.#size++;
+          if (this.#callback) {
+            this.#callback(node);
+            let parentNode = node.parent;
+            while (parentNode) {
+              this.#callback(parentNode);
+              parentNode = parentNode.parent;
+            }
+          }
           return node[direction];
         }
       }
@@ -410,9 +429,23 @@ export class BinarySearchTree<T> implements Iterable<T> {
     }
 
     this.#size--;
+    if (this.#callback) {
+      let parentNode = flaggedNode.parent;
+      while (parentNode) {
+        this.#callback(parentNode);
+        parentNode = parentNode.parent;
+      }
+    }
     return flaggedNode;
   }
 
+  /**
+   * Get the root node of the binary search tree.
+   */
+  getRoot(): BinarySearchNode<T> | null {
+    return this.#root;
+  }
+
   /**
    * Add a value to the binary search tree if it does not already exist in the
    * tree.

data_structures/binary_search_tree_test.ts

@@ -5,6 +5,7 @@ import {
   assertStrictEquals,
   assertThrows,
 } from "@std/assert";
+import { BinarySearchNode } from "./_binary_search_node.ts";
 import { BinarySearchTree } from "./binary_search_tree.ts";
 import { ascend, descend } from "./comparators.ts";
 
@@ -17,6 +18,14 @@ class MyMath {
 interface Container {
   id: number;
   values: number[];
+  st_size: number;
+}
+
+function callback(n: BinarySearchNode<Container>) {
+  let total_size = 1;
+  total_size += n.left?.value.st_size || 0;
+  total_size += n.right?.value.st_size || 0;
+  n.value.st_size = total_size;
 }
 
 Deno.test("BinarySearchTree throws if compare is not a function", () => {
@@ -271,28 +280,30 @@ Deno.test("BinarySearchTree contains objects", () => {
   const tree: BinarySearchTree<Container> = new BinarySearchTree((
     a: Container,
     b: Container,
-  ) => ascend(a.id, b.id));
+  ) => ascend(a.id, b.id), callback);
   const ids = [-10, 9, -1, 100, 1, 0, -100, 10, -9];
 
   for (const [i, id] of ids.entries()) {
-    const newContainer: Container = { id, values: [] };
+    const newContainer: Container = { id, values: [], st_size: 1 };
     assertEquals(tree.find(newContainer), null);
     assertEquals(tree.insert(newContainer), true);
     newContainer.values.push(i - 1, i, i + 1);
-    assertStrictEquals(tree.find({ id, values: [] }), newContainer);
+    assertStrictEquals(tree.find({ id, values: [], st_size: 1 }), newContainer);
     assertEquals(tree.size, i + 1);
     assertEquals(tree.isEmpty(), false);
+    assertEquals(tree.getRoot()?.value.st_size, i+1);
   }
+  assertEquals(tree.getRoot()?.value.st_size, ids.length);
   for (const [i, id] of ids.entries()) {
-    const newContainer: Container = { id, values: [] };
+    const newContainer: Container = { id, values: [], st_size: 1};
     assertEquals(
-      tree.find({ id } as Container),
-      { id, values: [i - 1, i, i + 1] },
+      tree.find({ id } as Container)?.id,
+      id,
     );
     assertEquals(tree.insert(newContainer), false);
     assertEquals(
-      tree.find({ id, values: [] }),
-      { id, values: [i - 1, i, i + 1] },
+      tree.find({ id, values: [], st_size: 1 })?.id,
+      id,
     );
     assertEquals(tree.size, ids.length);
     assertEquals(tree.isEmpty(), false);
@@ -310,18 +321,19 @@ Deno.test("BinarySearchTree contains objects", () => {
     assertEquals(tree.size, ids.length - i);
     assertEquals(tree.isEmpty(), false);
     assertEquals(
-      tree.find({ id, values: [] }),
-      { id, values: [i - 1, i, i + 1] },
+      tree.find({ id, values: [], st_size: 1 })?.id,
+      id,
     );
 
-    assertEquals(tree.remove({ id, values: [] }), true);
+    assertEquals(tree.remove({ id, values: [], st_size: 1 }), true);
+    assertEquals(tree.getRoot()?.value.st_size || 0, ids.length - i - 1);
     expected.splice(expected.indexOf(id), 1);
     assertEquals([...tree].map((container) => container.id), expected);
-    assertEquals(tree.find({ id, values: [] }), null);
+    assertEquals(tree.find({ id, values: [], st_size: 1 }), null);
 
-    assertEquals(tree.remove({ id, values: [] }), false);
+    assertEquals(tree.remove({ id, values: [], st_size: 1 }), false);
     assertEquals([...tree].map((container) => container.id), expected);
-    assertEquals(tree.find({ id, values: [] }), null);
+    assertEquals(tree.find({ id, values: [], st_size: 1 }), null);
   }
   assertEquals(tree.size, 0);
   assertEquals(tree.isEmpty(), true);

data_structures/red_black_tree.ts

@@ -3,6 +3,7 @@
 
 import { ascend } from "./comparators.ts";
 import { BinarySearchTree } from "./binary_search_tree.ts";
+import { BinarySearchNode } from "./_binary_search_node.ts";
 import { type Direction, RedBlackNode } from "./_red_black_node.ts";
 import { internals } from "./_binary_search_tree_internals.ts";
 
@@ -107,13 +108,18 @@ export class RedBlackTree<T> extends BinarySearchTree<T> {
    *
    * @param compare A custom comparison function for the values. The default comparison function sorts by ascending order.
    */
-  constructor(compare: (a: T, b: T) => number = ascend) {
+  constructor(compare: (a: T, b: T) => number = ascend, callback?: (node: BinarySearchNode<T>) => void) {
     if (typeof compare !== "function") {
       throw new TypeError(
         "Cannot construct a RedBlackTree: the 'compare' parameter is not a function, did you mean to call RedBlackTree.from?",
       );
     }
-    super(compare);
+    if (callback && typeof callback !== "function") {
+      throw new TypeError(
+        "Cannot construct a BinarySearchTree: the 'callback' parameter is not a function.",
+      );
+    }
+    super(compare, callback);
   }
 
   /**

data_structures/red_black_tree_test.ts

@@ -1,5 +1,7 @@
 // Copyright 2018-2025 the Deno authors. MIT license.
 import { assertEquals, assertStrictEquals, assertThrows } from "@std/assert";
+import { BinarySearchNode } from "./_binary_search_node.ts";
+import { RedBlackNode } from "./_red_black_node.ts";
 import { RedBlackTree } from "./red_black_tree.ts";
 import { ascend, descend } from "./comparators.ts";
 import { type Container, MyMath } from "./_test_utils.ts";
@@ -252,32 +254,40 @@ Deno.test("RedBlackTree works as exepcted with descend comparator", () => {
   }
 });
 
+function callback(n: BinarySearchNode<Container>) {
+  let total_size = 1;
+  total_size += n.left?.value.st_size || 0;
+  total_size += n.right?.value.st_size || 0;
+  n.value.st_size = total_size;
+}
+
 Deno.test("RedBlackTree works with object items", () => {
   const tree: RedBlackTree<Container> = new RedBlackTree((
     a: Container,
     b: Container,
-  ) => ascend(a.id, b.id));
+  ) => ascend(a.id, b.id), callback);
   const ids: number[] = [-10, 9, -1, 100, 1, 0, -100, 10, -9];
 
   for (const [i, id] of ids.entries()) {
-    const newContainer: Container = { id, values: [] };
+    const newContainer: Container = { id, values: [], st_size: 1 };
     assertEquals(tree.find(newContainer), null);
     assertEquals(tree.insert(newContainer), true);
     newContainer.values.push(i - 1, i, i + 1);
-    assertStrictEquals(tree.find({ id, values: [] }), newContainer);
+    assertStrictEquals(tree.find({ id, values: [], st_size: 1 }), newContainer);
     assertEquals(tree.size, i + 1);
+    assertEquals(tree.getRoot()?.value.st_size, i + 1);
     assertEquals(tree.isEmpty(), false);
   }
   for (const [i, id] of ids.entries()) {
-    const newContainer: Container = { id, values: [] };
+    const newContainer: Container = { id, values: [], st_size: 1 };
     assertEquals(
-      tree.find({ id } as Container),
-      { id, values: [i - 1, i, i + 1] },
+      tree.find({ id } as Container)?.id,
+      id,
     );
     assertEquals(tree.insert(newContainer), false);
     assertEquals(
-      tree.find({ id, values: [] }),
-      { id, values: [i - 1, i, i + 1] },
+      tree.find({ id, values: [], st_size: 1 })?.id,
+      id,
     );
     assertEquals(tree.size, ids.length);
     assertEquals(tree.isEmpty(), false);
@@ -295,18 +305,19 @@ Deno.test("RedBlackTree works with object items", () => {
     assertEquals(tree.size, ids.length - i);
     assertEquals(tree.isEmpty(), false);
     assertEquals(
-      tree.find({ id, values: [] }),
-      { id, values: [i - 1, i, i + 1] },
+      tree.find({ id, values: [], st_size: 1 })?.id,
+      id,
     );
 
-    assertEquals(tree.remove({ id, values: [] }), true);
+    assertEquals(tree.remove({ id, values: [], st_size: 1 }), true);
+    assertEquals(tree.getRoot()?.value.st_size || 0, ids.length - i - 1);
     expected.splice(expected.indexOf(id), 1);
     assertEquals([...tree].map((container) => container.id), expected);
-    assertEquals(tree.find({ id, values: [] }), null);
+    assertEquals(tree.find({ id, values: [], st_size: 1 }), null);
 
-    assertEquals(tree.remove({ id, values: [] }), false);
+    assertEquals(tree.remove({ id, values: [], st_size: 1 }), false);
     assertEquals([...tree].map((container) => container.id), expected);
-    assertEquals(tree.find({ id, values: [] }), null);
+    assertEquals(tree.find({ id, values: [], st_size: 1 }), null);
   }
   assertEquals(tree.size, 0);
   assertEquals(tree.isEmpty(), true);