diff --git a/src/data-structures/tree/avl-tree/AvlTree.js b/src/data-structures/tree/avl-tree/AvlTree.js index 51e6a05f..c26ef023 100644 --- 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 index fb1f3ee4..daf21413 100644 --- 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();