fix deletion bug and add rebalancing for avl delete

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<Number>, values: number[]) => {
+  for (const value of values) {
+    tree.insert(value)
+  }
+}
+
 describe('AVL tree', () => {
   let avlTree: AVLTree<Number>
 
@@ -9,39 +15,42 @@ 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)
     expect(avlTree.root?.right?.value).toBe(30)
   })
 
   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)
     expect(avlTree.root?.right?.value).toBe(60)
     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)
+  })
 })

src/avl-tree/index.ts

@@ -46,7 +46,7 @@ class AVLTree<Number> {
 
     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<Number> {
     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<Number> {
     value: Number
   ): AVLNode<Number> | 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<Number> | null): number {
+    if (!node) {
+      return 0
+    }
+
+    return (node.left?.height ?? 0) - (node.right?.height ?? 0)
+  }
+
   private rotateRight(x: AVLNode<Number>): AVLNode<Number> {
     const y = x.left as AVLNode<Number>
     const T2 = y.right as AVLNode<Number>

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', () => {

src/binary-tree/index.ts

@@ -96,7 +96,7 @@ export default class BinaryTree<Number> {
         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<Number> {
   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<Number> | 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)
+    }
+  }
 }