Simplify AVL tree node deletion.

trekhleb · trekhleb · commit afa49487678d · 2018-09-24T09:45:31.000+03:00
diff --git a/src/data-structures/tree/avl-tree/AvlTree.js b/src/data-structures/tree/avl-tree/AvlTree.js
@@ -16,54 +16,16 @@ export default class AvlTree extends BinarySearchTree {
     }
   }
 
-  remove(value) {
-    const nodeToRemove = this.root.find(value);
-
-    if (!nodeToRemove) {
-      throw new Error('Item not found in the tree');
-    }
-
-    // Recursively find target node, if found then delete and balance.
-    // return nodeToRemove.value;
-    this.root = this.removeRecv(this.root, nodeToRemove);
-  }
-
-  removeRecv(origin, node) {
-    let newOrigin = origin;
-    if (origin.value > node.value) {
-      // Recursively traversing from left
-      newOrigin.left = this.removeRecv(origin.left, node);
-    } else if (origin.value < node.value) {
-      // Recursively traversing from right
-      newOrigin.right = this.removeRecv(origin.right, node);
-    } else {
-      if (origin.left == null) {
-        // Forget right node
-        return origin.right;
-      }
-      if (origin.right == null) {
-        // Forget left node
-        return origin.left;
-      }
-
-      // Recursively find min node from left subtree
-      // more efficient traversing
-      const parent = Object.assign({}, origin);
-      newOrigin = parent.right.findMin();
-      newOrigin.right = this.deleteMin(parent.right);
-      newOrigin.left = parent.left;
-    }
-
-    // Balance and return root node
-    return this.balance(newOrigin);
-  }
-
   /**
    * @param {*} value
    * @return {boolean}
    */
   remove(value) {
-    throw new Error(`Can't remove ${value}. Remove method is not implemented yet`);
+    // Do standard BST removal.
+    super.remove(value);
+
+    // Balance the tree starting from the root node.
+    this.balance(this.root);
   }
 
   /**
@@ -90,8 +52,6 @@ export default class AvlTree extends BinarySearchTree {
         this.rotateRightLeft(node);
       }
     }
-    // Return the heap to avoid referenceError
-    return node;
   }
 
   /**
@@ -200,15 +160,4 @@ export default class AvlTree extends BinarySearchTree {
     // Attach rootNode to the left of rightNode.
     rightNode.setLeft(rootNode);
   }
-
-  deleteMin(node) {
-    // Forget right node if has value
-    if (node.left == null) return node.right;
-
-    const newNode = node;
-    newNode.left = this.deleteMin(node.left);
-
-    // Balance and return root node
-    return this.balance(newNode);
-  }
 }
diff --git a/src/data-structures/tree/avl-tree/__test__/AvlTRee.test.js b/src/data-structures/tree/avl-tree/__test__/AvlTRee.test.js
@@ -186,31 +186,6 @@ describe('AvlTree', () => {
     expect(tree.toString()).toBe('6,8,9,18,21,22,43');
   });
 
-  it('should keep balance after removal', () => {
-    const tree = new AvlTree();
-
-    tree.insert(1);
-    tree.insert(2);
-    tree.insert(3);
-    tree.insert(4);
-    tree.insert(5);
-    tree.insert(6);
-    tree.insert(7);
-    tree.insert(8);
-    tree.insert(9);
-
-    expect(tree.toString()).toBe('1,2,3,4,5,6,7,8,9');
-    expect(tree.root.height).toBe(3);
-
-    tree.remove(8);
-    tree.remove(9);
-
-    expect(tree.contains(8)).toBeFalsy();
-    expect(tree.contains(9)).toBeFalsy();
-    expect(tree.toString()).toBe('1,2,3,4,5,6,7');
-    expect(tree.root.height).toBe(2);
-  });
-
   it('should do left right rotation and keeping left right node safe', () => {
     const tree = new AvlTree();
 
@@ -255,13 +230,74 @@ describe('AvlTree', () => {
     expect(tree.root.height).toBe(3);
   });
 
-  it('should throw an error when trying to remove the node', () => {
-    const removeNodeAvlTree = () => {
-      const tree = new AvlTree();
+  it('should remove values from the tree with right-right rotation', () => {
+    const tree = new AvlTree();
+
+    tree.insert(10);
+    tree.insert(20);
+    tree.insert(30);
+    tree.insert(40);
+
+    expect(tree.toString()).toBe('10,20,30,40');
+
+    tree.remove(10);
+
+    expect(tree.toString()).toBe('20,30,40');
+    expect(tree.root.value).toBe(30);
+    expect(tree.root.left.value).toBe(20);
+    expect(tree.root.right.value).toBe(40);
+    expect(tree.root.balanceFactor).toBe(0);
+  });
+
+  it('should remove values from the tree with left-left rotation', () => {
+    const tree = new AvlTree();
+
+    tree.insert(10);
+    tree.insert(20);
+    tree.insert(30);
+    tree.insert(5);
+
+    expect(tree.toString()).toBe('5,10,20,30');
+
+    tree.remove(30);
+
+    expect(tree.toString()).toBe('5,10,20');
+    expect(tree.root.value).toBe(10);
+    expect(tree.root.left.value).toBe(5);
+    expect(tree.root.right.value).toBe(20);
+    expect(tree.root.balanceFactor).toBe(0);
+  });
+
+  it('should keep balance after removal', () => {
+    const tree = new AvlTree();
+
+    tree.insert(1);
+    tree.insert(2);
+    tree.insert(3);
+    tree.insert(4);
+    tree.insert(5);
+    tree.insert(6);
+    tree.insert(7);
+    tree.insert(8);
+    tree.insert(9);
+
+    expect(tree.toString()).toBe('1,2,3,4,5,6,7,8,9');
+    expect(tree.root.value).toBe(4);
+    expect(tree.root.height).toBe(3);
+    expect(tree.root.balanceFactor).toBe(-1);
+
+    tree.remove(8);
 
-      tree.remove(1);
-    };
+    expect(tree.root.value).toBe(4);
+    expect(tree.root.balanceFactor).toBe(-1);
+
+    tree.remove(9);
 
-    expect(removeNodeAvlTree).toThrowError();
+    expect(tree.contains(8)).toBeFalsy();
+    expect(tree.contains(9)).toBeFalsy();
+    expect(tree.toString()).toBe('1,2,3,4,5,6,7');
+    expect(tree.root.value).toBe(4);
+    expect(tree.root.height).toBe(2);
+    expect(tree.root.balanceFactor).toBe(0);
   });
 });
