feat: add callback functionality for binary search tree

data_structures/_binary_search_node.ts

@@ -1,13 +1,15 @@
 // Copyright 2018-2025 the Deno authors. MIT license.
 // This module is browser compatible.
 
+import type { BinarySearchTreeNode } from "./binary_search_tree_node.ts";
+
 export type Direction = "left" | "right";
 
-export class BinarySearchNode<T> {
-  left: BinarySearchNode<T> | null;
-  right: BinarySearchNode<T> | null;
-  parent: BinarySearchNode<T> | null;
-  value: T;
+export class BinarySearchNode<T> implements BinarySearchTreeNode<T> {
+  declare left: BinarySearchNode<T> | null;
+  declare right: BinarySearchNode<T> | null;
+  declare parent: BinarySearchNode<T> | null;
+  declare value: T;
 
   constructor(parent: BinarySearchNode<T> | null, value: T) {
     this.left = null;

data_structures/binary_search_tree.ts

@@ -2,6 +2,7 @@
 // This module is browser compatible.
 
 import { ascend } from "./comparators.ts";
+import type { BinarySearchTreeNode } from "./binary_search_tree_node.ts";
 import { BinarySearchNode } from "./_binary_search_node.ts";
 import { internals } from "./_binary_search_tree_internals.ts";
 
@@ -93,6 +94,7 @@ type Direction = "left" | "right";
 export class BinarySearchTree<T> implements Iterable<T> {
   #root: BinarySearchNode<T> | null = null;
   #size = 0;
+  #callback: ((node: BinarySearchTreeNode<T>) => void) | null = null;
   #compare: (a: T, b: T) => number;
 
   /**
@@ -103,13 +105,25 @@ 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.
+   * @param callback An optional callback function that is called whenever a change
+   * is made in the subtree of a node. This is guaranteed to be called in order from
+   * leaves to the root.
    */
-  constructor(compare: (a: T, b: T) => number = ascend) {
+  constructor(
+    compare: (a: T, b: T) => number = ascend,
+    callback?: (node: BinarySearchTreeNode<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 +367,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 +392,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 +436,37 @@ 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.
+   *
+   * @example Getting the root node of the tree
+   * ```ts
+   * import { BinarySearchTree } from "@std/data-structures";
+   * import { assertEquals } from "@std/assert";
+   *
+   * const tree = new BinarySearchTree<number>();
+   *
+   * assertEquals(tree.insert(42), true);
+   * let root = tree.getRoot();
+   * assertEquals(root?.value, 42);
+   * ```
+   *
+   * @returns A reference to the root node of the binary search tree, or null if the tree is empty.
+   */
+  getRoot(): BinarySearchTreeNode<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_node.ts

@@ -0,0 +1,108 @@
+// Copyright 2018-2025 the Deno authors. MIT license.
+// This module is browser compatible.
+
+/**
+ * A generic Binary Search Tree (BST) node interface.
+ * It is implemented by the internal node classes of the binary search tree and
+ * red black tree.
+ *
+ * @example Getting a reference to a BinarySearchTreeNode<T>.
+ * ```ts
+ * import { BinarySearchTree<T> } from "@std/data-structures";
+ * import { assertEquals } from "@std/assert";
+ *
+ * const tree = new BinarySearchTree<number>();
+ * 
+ * assertEquals(tree.insert(42), true);
+ * 
+ * const root = tree.getRoot();
+ * 
+ * assertEquals(node.value, 42);
+ * assertEquals(node.left, null);
+ * assertEquals(node.right, null);
+ * assertEquals(node.parent, null);
+ * ```
+ *
+ * @typeparam T The type of the values stored in the binary tree.
+ */
+export interface BinarySearchTreeNode<T> {
+  /**
+   * The left child node, or null if there is no left child.
+   *
+   * @example Checking the left child of a node in a binary search tree.
+   * ```ts
+   * import { BinarySearchTree } from "@std/data-structures";
+   * import { assertEquals } from "@std/assert";
+   *
+   * const tree = new BinarySearchTree<number>();
+   *
+   * assertEquals(tree.insert(42), true);
+   * assertEquals(tree.insert(21), true);
+   *
+   * const root = tree.getRoot();
+   * const leftChild = root?.left;
+   *
+   * assertEquals(leftChild?.value, 21);
+   * ```
+   */
+  left: BinarySearchTreeNode<T> | null;
+
+  /**
+   * The right child node, or null if there is no right child.
+   *
+   * @example Checking the right child of a node in a binary search tree.
+   * ```ts
+   * import { BinarySearchTree } from "@std/data-structures";
+   * import { assertEquals } from "@std/assert";
+   *
+   * const tree = new BinarySearchTree<number>();
+   *
+   * assertEquals(tree.insert(21), true);
+   * assertEquals(tree.insert(42), true);
+   *
+   * const root = tree.getRoot();
+   * const leftChild = root?.left;
+   *
+   * assertEquals(leftChild?.value, 42);
+   * ```
+   */
+  right: BinarySearchTreeNode<T> | null;
+
+  /**
+   * The parent of this node, or null if there is no parent.
+   *
+   * @example Checking the parent of a node in a binary search tree.
+   * ```ts
+   * import { BinarySearchTree } from "@std/data-structures";
+   * import { assertEquals } from "@std/assert";
+   *
+   * const tree = new BinarySearchTree<number>();
+   *
+   * assertEquals(tree.insert(42), true);
+   * assertEquals(tree.insert(21), true);
+   *
+   * const root = tree.getRoot();
+   * const leftChild = root?.left;
+   *
+   * assertEquals(leftChild?.parent?.value, 42);
+   * ```
+   */
+  parent: BinarySearchTreeNode<T> | null;
+
+  /**
+   * The value stored at this node.
+   *
+   * @example Accessing the value of a node in a binary search tree.
+   * ```ts
+   * import { BinarySearchTree } from "@std/data-structures";
+   * import { assertEquals } from "@std/assert";
+   *
+   * const tree = new BinarySearchTree<number>();
+   * assertEquals(tree.insert(42), true);
+   *
+   * const root = tree.getRoot();
+   * assertEquals(root?.value, 42);
+   * ```
+   */
+  value: T;
+}

data_structures/binary_search_tree_test.ts

@@ -5,6 +5,7 @@ import {
   assertStrictEquals,
   assertThrows,
 } from "@std/assert";
+import type { BinarySearchTreeNode } from "./binary_search_tree_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[];
+  stSize: number;
+}
+
+function callback(n: BinarySearchTreeNode<Container>) {
+  let totalSize = 1;
+  totalSize += n.left?.value.stSize || 0;
+  totalSize += n.right?.value.stSize || 0;
+  n.value.stSize = totalSize;
 }
 
 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: [], stSize: 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: [], stSize: 1 }), newContainer);
     assertEquals(tree.size, i + 1);
     assertEquals(tree.isEmpty(), false);
+    assertEquals(tree.getRoot()?.value.stSize, i + 1);
   }
-  for (const [i, id] of ids.entries()) {
-    const newContainer: Container = { id, values: [] };
+  assertEquals(tree.getRoot()?.value.stSize, ids.length);
+  for (const [_i, id] of ids.entries()) {
+    const newContainer: Container = { id, values: [], stSize: 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: [], stSize: 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: [], stSize: 1 })?.id,
+      id,
     );
 
-    assertEquals(tree.remove({ id, values: [] }), true);
+    assertEquals(tree.remove({ id, values: [], stSize: 1 }), true);
+    assertEquals(tree.getRoot()?.value.stSize || 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: [], stSize: 1 }), null);
 
-    assertEquals(tree.remove({ id, values: [] }), false);
+    assertEquals(tree.remove({ id, values: [], stSize: 1 }), false);
     assertEquals([...tree].map((container) => container.id), expected);
-    assertEquals(tree.find({ id, values: [] }), null);
+    assertEquals(tree.find({ id, values: [], stSize: 1 }), null);
   }
   assertEquals(tree.size, 0);
   assertEquals(tree.isEmpty(), true);

data_structures/mod.ts

@@ -27,5 +27,6 @@
 
 export * from "./binary_heap.ts";
 export * from "./binary_search_tree.ts";
+export * from "./binary_search_tree_node.ts";
 export * from "./comparators.ts";
 export * from "./red_black_tree.ts";

data_structures/red_black_tree.ts

@@ -3,6 +3,7 @@
 
 import { ascend } from "./comparators.ts";
 import { BinarySearchTree } from "./binary_search_tree.ts";
+import type { BinarySearchTreeNode } from "./binary_search_tree_node.ts";
 import { type Direction, RedBlackNode } from "./_red_black_node.ts";
 import { internals } from "./_binary_search_tree_internals.ts";
 
@@ -106,14 +107,23 @@ export class RedBlackTree<T> extends BinarySearchTree<T> {
    * Construct an empty red-black tree.
    *
    * @param compare A custom comparison function for the values. The default comparison function sorts by ascending order.
+   * @param callback An optional callback function that is called whenever a change is made in the subtree of a node. This is guaranteed to be called in order from leaves to the root.
    */
-  constructor(compare: (a: T, b: T) => number = ascend) {
+  constructor(
+    compare: (a: T, b: T) => number = ascend,
+    callback?: (node: BinarySearchTreeNode<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 RedBlackTree: the 'callback' parameter is not a function",
+      );
+    }
+    super(compare, callback);
   }
 
   /**