Add remove method (trekhleb#33)

Bladefidz · trekhleb · commit 04e533e4b511 · 2018-09-24T09:12:38.000+03:00
Remove node in AvlTree with auto balancing. Fix issue: trekhleb#13
diff --git a/src/data-structures/tree/avl-tree/AvlTree.js b/src/data-structures/tree/avl-tree/AvlTree.js
@@ -16,6 +16,48 @@ 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}
@@ -48,6 +90,8 @@ export default class AvlTree extends BinarySearchTree {
         this.rotateRightLeft(node);
       }
     }
+    // Return the heap to avoid referenceError
+    return node;
   }
 
   /**
@@ -156,4 +200,15 @@ 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,6 +186,31 @@ 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();
 
