From 5fe184c1527ae746a8e8f86e647749e7e365d895 Mon Sep 17 00:00:00 2001 From: Joe Carpenito Date: Mon, 13 Jan 2025 10:38:08 -0500 Subject: [PATCH] fix deletion bug and add rebalancing for avl delete --- src/avl-tree/index.spec.ts | 51 ++++++++++++++++++++-------------- src/avl-tree/index.ts | 52 ++++++++++++++++++++++++++--------- src/binary-tree/index.spec.ts | 2 ++ src/binary-tree/index.ts | 26 +++++++++++++++++- 4 files changed, 96 insertions(+), 35 deletions(-) diff --git a/src/avl-tree/index.spec.ts b/src/avl-tree/index.spec.ts index c777cee..e7e4c88 100644 --- a/src/avl-tree/index.spec.ts +++ b/src/avl-tree/index.spec.ts @@ -1,6 +1,12 @@ import { beforeEach, describe, expect, test } from 'bun:test' import AVLTree from '.' +const populateTree = (tree: AVLTree, values: number[]) => { + for (const value of values) { + tree.insert(value) + } +} + describe('AVL tree', () => { let avlTree: AVLTree @@ -9,9 +15,7 @@ describe('AVL tree', () => { }) test('should create a binary tree that is balanced', () => { - avlTree.insert(30) - avlTree.insert(20) - avlTree.insert(10) + populateTree(avlTree, [30, 20, 10]) expect(avlTree.root?.value).toBe(20) expect(avlTree.root?.left?.value).toBe(10) @@ -19,24 +23,10 @@ describe('AVL tree', () => { }) test('should insert multiple values into the tree', () => { - avlTree.insert(30) - avlTree.insert(40) - avlTree.insert(50) - avlTree.insert(20) - avlTree.insert(10) - avlTree.insert(5) - avlTree.insert(15) - avlTree.insert(25) - avlTree.insert(35) - avlTree.insert(45) - avlTree.insert(55) - avlTree.insert(60) - avlTree.insert(70) - avlTree.insert(65) - avlTree.insert(75) - avlTree.insert(80) - avlTree.insert(90) - avlTree.insert(95) + populateTree( + avlTree, + [30, 40, 50, 20, 10, 5, 15, 25, 35, 45, 55, 60, 70, 65, 75, 80, 90, 95] + ) expect(avlTree.root?.value).toBe(40) expect(avlTree.root?.left?.value).toBe(20) @@ -44,4 +34,23 @@ describe('AVL tree', () => { expect(avlTree.root?.left?.left?.left?.value).toBe(5) expect(avlTree.root?.right?.right?.right?.right?.value).toBe(95) }) + + test('should remove nodes from the tree and successfully rebalance', () => { + populateTree(avlTree, [30, 20, 10]) + + avlTree.remove(30) + + expect(avlTree.root?.value).toBe(20) + expect(avlTree.root?.left?.value).toBe(10) + + avlTree.clear() + + populateTree(avlTree, [80, 70, 60, 50, 40, 30, 20, 10]) + + avlTree.remove(30) + + expect(avlTree.root?.value).toBe(50) + expect(avlTree.root?.left?.value).toBe(20) + expect(avlTree.root?.left?.left?.value).toBe(10) + }) }) diff --git a/src/avl-tree/index.ts b/src/avl-tree/index.ts index cd8d90f..bf82ed6 100644 --- a/src/avl-tree/index.ts +++ b/src/avl-tree/index.ts @@ -46,7 +46,7 @@ class AVLTree { node.height = Math.max(node.left?.height ?? 0, node.right?.height ?? 0) + 1 - const balance = (node.left?.height ?? 0) - (node.right?.height ?? 0) + const balance = this.getBalance(node) if (balance > 1 && node.left && node.value > node.left.value) { return this.rotateRight(node) @@ -76,14 +76,11 @@ class AVLTree { if (!node) { return null } else if (node.value === value) { - // we found the node with the value we are looking for return node } else if (value < node.value) { - // traverse left return this.findRecursion(node.left, value) } - // traverse right return this.findRecursion(node.right, value) } @@ -92,37 +89,66 @@ class AVLTree { value: Number ): AVLNode | null => { if (!node) { - return null + return node } else if (value < node.value) { - // if the value we want to remove is less than the current node's value, - // we recurse to the left node.left = this.removeRecursion(node.left, value) } else if (value > node.value) { - // if the value we want to remove is greater than the current node's value, - // we recurse to the right node.right = this.removeRecursion(node.right, value) } else { if (node.left === null && node.right === null) { // if the node has no children, we can just remove it node = null } else if (node.left == null) { - // if the node has only a right child, we can replace the node with its right child node = node.right } else if (node.right == null) { - // if the node has only a left child, we can replace the node with its left child node = node.left } else { - // if the node has two children, we find the minimum value in the right subtree const minRight = node.right.findMin() node.value = minRight.value - node.right = this.removeRecursion(node.right, value) + node.right = this.removeRecursion(node.right, minRight.value) } } + if (!node) { + return null + } + + // rebalancing for deletion is different from insertion as + // we need to check the balance of the node after deletion + node.height = Math.max(node.left?.height ?? 0, node.right?.height ?? 0) + 1 + + const balance = this.getBalance(node) + + if (balance > 1 && this.getBalance(node.left) >= 0) { + return this.rotateRight(node) + } + + if (balance < -1 && this.getBalance(node.right) <= 0) { + return this.rotateLeft(node) + } + + if (balance > 1 && node.left && this.getBalance(node.left) < 0) { + node.left = this.rotateLeft(node.left) + return this.rotateRight(node) + } + + if (balance < -1 && node.right && this.getBalance(node.right) > 0) { + node.right = this.rotateRight(node.right) + return this.rotateLeft(node) + } + return node } + private getBalance(node: AVLNode | null): number { + if (!node) { + return 0 + } + + return (node.left?.height ?? 0) - (node.right?.height ?? 0) + } + private rotateRight(x: AVLNode): AVLNode { const y = x.left as AVLNode const T2 = y.right as AVLNode diff --git a/src/binary-tree/index.spec.ts b/src/binary-tree/index.spec.ts index fc444d8..e268aa4 100644 --- a/src/binary-tree/index.spec.ts +++ b/src/binary-tree/index.spec.ts @@ -94,6 +94,8 @@ describe('BinaryTree', () => { tree.remove(15) expect(tree.root?.right?.value).toBe(18) + expect(tree.root?.right?.right?.value).toBe(20) + expect(tree.root?.right?.right?.left?.value).toBeUndefined() // <-- was 18 }) test('should handle removing a non-existent value gracefully', () => { diff --git a/src/binary-tree/index.ts b/src/binary-tree/index.ts index 583ec20..dfbbee5 100644 --- a/src/binary-tree/index.ts +++ b/src/binary-tree/index.ts @@ -96,7 +96,7 @@ export default class BinaryTree { const minRight = node.right.findMin() node.value = minRight.value - node.right = this.removeRecursion(node.right, value) + node.right = this.removeRecursion(node.right, minRight.value) } } @@ -155,4 +155,28 @@ export default class BinaryTree { public remove(value: Number): void { this.removeRecursion(this._root, value) } + + /** + * Prints the binary tree in a console-friendly format that clearly illustrates the tree structure. + * @returns {void} + */ + public print(): void { + const printNode = ( + node: Node | null, + prefix: string = '', + isLeft: boolean = true + ) => { + if (node !== null) { + console.log(prefix + (isLeft ? '├── ' : '└── ') + node.value) + printNode(node.left, prefix + (isLeft ? '│ ' : ' '), true) + printNode(node.right, prefix + (isLeft ? '│ ' : ' '), false) + } + } + + if (this._root === null) { + console.log('Tree is empty') + } else { + printNode(this._root, '', true) + } + } }