Skip to content

Commit d9443d5

Browse files
committed
feat: add callback functionality for binary search tree
Signed-off-by: Abhinav Tamaskar <[email protected]>
1 parent 2258bf2 commit d9443d5

File tree

5 files changed

+94
-31
lines changed

5 files changed

+94
-31
lines changed

data_structures/_test_utils.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,4 +8,5 @@ export class MyMath {
88
export interface Container {
99
id: number;
1010
values: number[];
11+
st_size: number;
1112
}

data_structures/binary_search_tree.ts

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,7 @@ type Direction = "left" | "right";
9393
export class BinarySearchTree<T> implements Iterable<T> {
9494
#root: BinarySearchNode<T> | null = null;
9595
#size = 0;
96+
#callback: ((node: BinarySearchNode<T>) => void) | null = null;
9697
#compare: (a: T, b: T) => number;
9798

9899
/**
@@ -104,12 +105,18 @@ export class BinarySearchTree<T> implements Iterable<T> {
104105
* @param compare A custom comparison function to sort the values in the tree.
105106
* By default, the values are sorted in ascending order.
106107
*/
107-
constructor(compare: (a: T, b: T) => number = ascend) {
108+
constructor(compare: (a: T, b: T) => number = ascend, callback?: (node: BinarySearchNode<T>) => void) {
108109
if (typeof compare !== "function") {
109110
throw new TypeError(
110111
"Cannot construct a BinarySearchTree: the 'compare' parameter is not a function, did you mean to call BinarySearchTree.from?",
111112
);
112113
}
114+
if (callback && typeof callback !== "function") {
115+
throw new TypeError(
116+
"Cannot construct a BinarySearchTree: the 'callback' parameter is not a function.",
117+
);
118+
}
119+
this.#callback = callback || null;
113120
this.#compare = compare;
114121
}
115122

@@ -353,6 +360,10 @@ export class BinarySearchTree<T> implements Iterable<T> {
353360
}
354361
replacement[direction] = node;
355362
node.parent = replacement;
363+
if (this.#callback) {
364+
this.#callback(node);
365+
this.#callback(replacement);
366+
}
356367
}
357368

358369
#insertNode(
@@ -374,6 +385,14 @@ export class BinarySearchTree<T> implements Iterable<T> {
374385
} else {
375386
node[direction] = new Node(node, value);
376387
this.#size++;
388+
if (this.#callback) {
389+
this.#callback(node);
390+
let parentNode = node.parent;
391+
while (parentNode) {
392+
this.#callback(parentNode);
393+
parentNode = parentNode.parent;
394+
}
395+
}
377396
return node[direction];
378397
}
379398
}
@@ -410,9 +429,23 @@ export class BinarySearchTree<T> implements Iterable<T> {
410429
}
411430

412431
this.#size--;
432+
if (this.#callback) {
433+
let parentNode = flaggedNode.parent;
434+
while (parentNode) {
435+
this.#callback(parentNode);
436+
parentNode = parentNode.parent;
437+
}
438+
}
413439
return flaggedNode;
414440
}
415441

442+
/**
443+
* Get the root node of the binary search tree.
444+
*/
445+
getRoot(): BinarySearchNode<T> | null {
446+
return this.#root;
447+
}
448+
416449
/**
417450
* Add a value to the binary search tree if it does not already exist in the
418451
* tree.

data_structures/binary_search_tree_test.ts

Lines changed: 26 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import {
55
assertStrictEquals,
66
assertThrows,
77
} from "@std/assert";
8+
import { BinarySearchNode } from "./_binary_search_node.ts";
89
import { BinarySearchTree } from "./binary_search_tree.ts";
910
import { ascend, descend } from "./comparators.ts";
1011

@@ -17,6 +18,14 @@ class MyMath {
1718
interface Container {
1819
id: number;
1920
values: number[];
21+
st_size: number;
22+
}
23+
24+
function callback(n: BinarySearchNode<Container>) {
25+
let total_size = 1;
26+
total_size += n.left?.value.st_size || 0;
27+
total_size += n.right?.value.st_size || 0;
28+
n.value.st_size = total_size;
2029
}
2130

2231
Deno.test("BinarySearchTree throws if compare is not a function", () => {
@@ -271,28 +280,30 @@ Deno.test("BinarySearchTree contains objects", () => {
271280
const tree: BinarySearchTree<Container> = new BinarySearchTree((
272281
a: Container,
273282
b: Container,
274-
) => ascend(a.id, b.id));
283+
) => ascend(a.id, b.id), callback);
275284
const ids = [-10, 9, -1, 100, 1, 0, -100, 10, -9];
276285

277286
for (const [i, id] of ids.entries()) {
278-
const newContainer: Container = { id, values: [] };
287+
const newContainer: Container = { id, values: [], st_size: 1 };
279288
assertEquals(tree.find(newContainer), null);
280289
assertEquals(tree.insert(newContainer), true);
281290
newContainer.values.push(i - 1, i, i + 1);
282-
assertStrictEquals(tree.find({ id, values: [] }), newContainer);
291+
assertStrictEquals(tree.find({ id, values: [], st_size: 1 }), newContainer);
283292
assertEquals(tree.size, i + 1);
284293
assertEquals(tree.isEmpty(), false);
294+
assertEquals(tree.getRoot()?.value.st_size, i+1);
285295
}
296+
assertEquals(tree.getRoot()?.value.st_size, ids.length);
286297
for (const [i, id] of ids.entries()) {
287-
const newContainer: Container = { id, values: [] };
298+
const newContainer: Container = { id, values: [], st_size: 1};
288299
assertEquals(
289-
tree.find({ id } as Container),
290-
{ id, values: [i - 1, i, i + 1] },
300+
tree.find({ id } as Container)?.id,
301+
id,
291302
);
292303
assertEquals(tree.insert(newContainer), false);
293304
assertEquals(
294-
tree.find({ id, values: [] }),
295-
{ id, values: [i - 1, i, i + 1] },
305+
tree.find({ id, values: [], st_size: 1 })?.id,
306+
id,
296307
);
297308
assertEquals(tree.size, ids.length);
298309
assertEquals(tree.isEmpty(), false);
@@ -310,18 +321,19 @@ Deno.test("BinarySearchTree contains objects", () => {
310321
assertEquals(tree.size, ids.length - i);
311322
assertEquals(tree.isEmpty(), false);
312323
assertEquals(
313-
tree.find({ id, values: [] }),
314-
{ id, values: [i - 1, i, i + 1] },
324+
tree.find({ id, values: [], st_size: 1 })?.id,
325+
id,
315326
);
316327

317-
assertEquals(tree.remove({ id, values: [] }), true);
328+
assertEquals(tree.remove({ id, values: [], st_size: 1 }), true);
329+
assertEquals(tree.getRoot()?.value.st_size || 0, ids.length - i - 1);
318330
expected.splice(expected.indexOf(id), 1);
319331
assertEquals([...tree].map((container) => container.id), expected);
320-
assertEquals(tree.find({ id, values: [] }), null);
332+
assertEquals(tree.find({ id, values: [], st_size: 1 }), null);
321333

322-
assertEquals(tree.remove({ id, values: [] }), false);
334+
assertEquals(tree.remove({ id, values: [], st_size: 1 }), false);
323335
assertEquals([...tree].map((container) => container.id), expected);
324-
assertEquals(tree.find({ id, values: [] }), null);
336+
assertEquals(tree.find({ id, values: [], st_size: 1 }), null);
325337
}
326338
assertEquals(tree.size, 0);
327339
assertEquals(tree.isEmpty(), true);

data_structures/red_black_tree.ts

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
import { ascend } from "./comparators.ts";
55
import { BinarySearchTree } from "./binary_search_tree.ts";
6+
import { BinarySearchNode } from "./_binary_search_node.ts";
67
import { type Direction, RedBlackNode } from "./_red_black_node.ts";
78
import { internals } from "./_binary_search_tree_internals.ts";
89

@@ -107,13 +108,18 @@ export class RedBlackTree<T> extends BinarySearchTree<T> {
107108
*
108109
* @param compare A custom comparison function for the values. The default comparison function sorts by ascending order.
109110
*/
110-
constructor(compare: (a: T, b: T) => number = ascend) {
111+
constructor(compare: (a: T, b: T) => number = ascend, callback?: (node: BinarySearchNode<T>) => void) {
111112
if (typeof compare !== "function") {
112113
throw new TypeError(
113114
"Cannot construct a RedBlackTree: the 'compare' parameter is not a function, did you mean to call RedBlackTree.from?",
114115
);
115116
}
116-
super(compare);
117+
if (callback && typeof callback !== "function") {
118+
throw new TypeError(
119+
"Cannot construct a BinarySearchTree: the 'callback' parameter is not a function.",
120+
);
121+
}
122+
super(compare, callback);
117123
}
118124

119125
/**

data_structures/red_black_tree_test.ts

Lines changed: 25 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
// Copyright 2018-2025 the Deno authors. MIT license.
22
import { assertEquals, assertStrictEquals, assertThrows } from "@std/assert";
3+
import { BinarySearchNode } from "./_binary_search_node.ts";
4+
import { RedBlackNode } from "./_red_black_node.ts";
35
import { RedBlackTree } from "./red_black_tree.ts";
46
import { ascend, descend } from "./comparators.ts";
57
import { type Container, MyMath } from "./_test_utils.ts";
@@ -252,32 +254,40 @@ Deno.test("RedBlackTree works as exepcted with descend comparator", () => {
252254
}
253255
});
254256

257+
function callback(n: BinarySearchNode<Container>) {
258+
let total_size = 1;
259+
total_size += n.left?.value.st_size || 0;
260+
total_size += n.right?.value.st_size || 0;
261+
n.value.st_size = total_size;
262+
}
263+
255264
Deno.test("RedBlackTree works with object items", () => {
256265
const tree: RedBlackTree<Container> = new RedBlackTree((
257266
a: Container,
258267
b: Container,
259-
) => ascend(a.id, b.id));
268+
) => ascend(a.id, b.id), callback);
260269
const ids: number[] = [-10, 9, -1, 100, 1, 0, -100, 10, -9];
261270

262271
for (const [i, id] of ids.entries()) {
263-
const newContainer: Container = { id, values: [] };
272+
const newContainer: Container = { id, values: [], st_size: 1 };
264273
assertEquals(tree.find(newContainer), null);
265274
assertEquals(tree.insert(newContainer), true);
266275
newContainer.values.push(i - 1, i, i + 1);
267-
assertStrictEquals(tree.find({ id, values: [] }), newContainer);
276+
assertStrictEquals(tree.find({ id, values: [], st_size: 1 }), newContainer);
268277
assertEquals(tree.size, i + 1);
278+
assertEquals(tree.getRoot()?.value.st_size, i + 1);
269279
assertEquals(tree.isEmpty(), false);
270280
}
271281
for (const [i, id] of ids.entries()) {
272-
const newContainer: Container = { id, values: [] };
282+
const newContainer: Container = { id, values: [], st_size: 1 };
273283
assertEquals(
274-
tree.find({ id } as Container),
275-
{ id, values: [i - 1, i, i + 1] },
284+
tree.find({ id } as Container)?.id,
285+
id,
276286
);
277287
assertEquals(tree.insert(newContainer), false);
278288
assertEquals(
279-
tree.find({ id, values: [] }),
280-
{ id, values: [i - 1, i, i + 1] },
289+
tree.find({ id, values: [], st_size: 1 })?.id,
290+
id,
281291
);
282292
assertEquals(tree.size, ids.length);
283293
assertEquals(tree.isEmpty(), false);
@@ -295,18 +305,19 @@ Deno.test("RedBlackTree works with object items", () => {
295305
assertEquals(tree.size, ids.length - i);
296306
assertEquals(tree.isEmpty(), false);
297307
assertEquals(
298-
tree.find({ id, values: [] }),
299-
{ id, values: [i - 1, i, i + 1] },
308+
tree.find({ id, values: [], st_size: 1 })?.id,
309+
id,
300310
);
301311

302-
assertEquals(tree.remove({ id, values: [] }), true);
312+
assertEquals(tree.remove({ id, values: [], st_size: 1 }), true);
313+
assertEquals(tree.getRoot()?.value.st_size || 0, ids.length - i - 1);
303314
expected.splice(expected.indexOf(id), 1);
304315
assertEquals([...tree].map((container) => container.id), expected);
305-
assertEquals(tree.find({ id, values: [] }), null);
316+
assertEquals(tree.find({ id, values: [], st_size: 1 }), null);
306317

307-
assertEquals(tree.remove({ id, values: [] }), false);
318+
assertEquals(tree.remove({ id, values: [], st_size: 1 }), false);
308319
assertEquals([...tree].map((container) => container.id), expected);
309-
assertEquals(tree.find({ id, values: [] }), null);
320+
assertEquals(tree.find({ id, values: [], st_size: 1 }), null);
310321
}
311322
assertEquals(tree.size, 0);
312323
assertEquals(tree.isEmpty(), true);

0 commit comments

Comments
 (0)