diff --git a/src/data-structures/tree/avl-tree/AvlTree.js b/src/data-structures/tree/avl-tree/AvlTree.js new file mode 100644 index 00000000..f3005892 --- /dev/null +++ b/src/data-structures/tree/avl-tree/AvlTree.js @@ -0,0 +1,122 @@ +import BinarySearchTree from '../binary-search-tree/BinarySearchTree'; + +export default class AvlTree extends BinarySearchTree { + insert(value) { + // Do the normal BST insert. + super.insert(value); + + // Let's move up to the root and check balance factors along the way. + let currentNode = this.root.find(value); + while (currentNode) { + this.balance(currentNode); + currentNode = currentNode.parent; + } + } + + balance(node) { + // If balance factor is not OK then try to balance the node. + if (node.balanceFactor > 1) { + // Left rotation. + if (node.left.balanceFactor > 0) { + // Left-Left rotation + this.rotateLeftLeft(node); + } else if (node.left.balanceFactor < 0) { + // Left-Right rotation. + this.rotateLeftRight(node); + } + } else if (node.balanceFactor < -1) { + // Right rotation. + if (node.right.balanceFactor < 0) { + // Right-Right rotation + this.rotateRightRight(node); + } else if (node.right.balanceFactor > 0) { + // Right-Left rotation. + this.rotateRightLeft(node); + } + } + } + + rotateLeftLeft(rootNode) { + // Detach left node from root node. + const leftNode = rootNode.left; + rootNode.setLeft(null); + + // Make left node to be a child of rootNode's parent. + if (rootNode.parent) { + rootNode.parent.setLeft(leftNode); + } else if (rootNode === this.root) { + // If root node is root then make left node to be a new root. + this.root = leftNode; + } + + // If left node has a right child then detach it and + // attach it as a left child for rootNode. + if (leftNode.right) { + rootNode.setLeft(leftNode.right); + } + + // Attach rootNode to the right of leftNode. + leftNode.setRight(rootNode); + } + + rotateLeftRight(rootNode) { + // Detach left node from rootNode since it is going to be replaced. + const leftNode = rootNode.left; + rootNode.setLeft(null); + + // Detach right node from leftNode. + const leftRightNode = leftNode.right; + leftNode.setRight(null); + + // Attach leftRightNode to the rootNode. + rootNode.setLeft(leftRightNode); + + // Attach leftNode as left node for leftRight node. + leftRightNode.setLeft(leftNode); + + // Do left-left rotation. + this.rotateLeftLeft(rootNode); + } + + rotateRightLeft(rootNode) { + // Detach right node from rootNode since it is going to be replaced. + const rightNode = rootNode.right; + rootNode.setRight(null); + + // Detach left node from rightNode. + const rightLeftNode = rightNode.left; + rightNode.setLeft(null); + + // Attach rightLeftNode to the rootNode. + rootNode.setRight(rightLeftNode); + + // Attach rightNode as right node for rightLeft node. + rightLeftNode.setRight(rightNode); + + // Do right-right rotation. + this.rotateRightRight(rootNode); + } + + rotateRightRight(rootNode) { + // Detach right node from root node. + const rightNode = rootNode.right; + rootNode.setRight(null); + + // Make right node to be a child of rootNode's parent. + if (rootNode.parent) { + rootNode.parent.setRight(rightNode); + } else if (rootNode === this.root) { + // If root node is root then make right node to be a new root. + this.root = rightNode; + } + + // If right node has a left child then detach it and + // attach it as a right child for rootNode. + if (rightNode.left) { + rootNode.setRight(rightNode.left); + } + + // Attach rootNode to the left of rightNode. + rightNode.setLeft(rootNode); + } +} diff --git a/src/data-structures/tree/avl-tree/__test__/AvlTRee.test.js b/src/data-structures/tree/avl-tree/__test__/AvlTRee.test.js new file mode 100644 index 00000000..3411407d --- /dev/null +++ b/src/data-structures/tree/avl-tree/__test__/AvlTRee.test.js @@ -0,0 +1,165 @@ +import AvlTree from '../AvlTree'; + +describe('AvlTree', () => { + it('should do simple left-left rotation', () => { + const tree = new AvlTree(); + + tree.insert(4); + tree.insert(3); + tree.insert(2); + + expect(tree.toString()).toBe('2,3,4'); + expect(tree.root.value).toBe(3); + expect(tree.root.height).toBe(1); + + tree.insert(1); + + expect(tree.toString()).toBe('1,2,3,4'); + expect(tree.root.value).toBe(3); + expect(tree.root.height).toBe(2); + + tree.insert(0); + + expect(tree.toString()).toBe('0,1,2,3,4'); + expect(tree.root.value).toBe(3); + expect(tree.root.left.value).toBe(1); + expect(tree.root.height).toBe(2); + }); + + it('should do complex left-left rotation', () => { + const tree = new AvlTree(); + + tree.insert(30); + tree.insert(20); + tree.insert(40); + tree.insert(10); + + expect(tree.root.value).toBe(30); + expect(tree.root.height).toBe(2); + expect(tree.toString()).toBe('10,20,30,40'); + + tree.insert(25); + expect(tree.root.value).toBe(30); + expect(tree.root.height).toBe(2); + expect(tree.toString()).toBe('10,20,25,30,40'); + + tree.insert(5); + expect(tree.root.value).toBe(20); + expect(tree.root.height).toBe(2); + expect(tree.toString()).toBe('5,10,20,25,30,40'); + }); + + it('should do simple right-right rotation', () => { + const tree = new AvlTree(); + + tree.insert(2); + tree.insert(3); + tree.insert(4); + + expect(tree.toString()).toBe('2,3,4'); + expect(tree.root.value).toBe(3); + expect(tree.root.height).toBe(1); + + tree.insert(5); + + expect(tree.toString()).toBe('2,3,4,5'); + expect(tree.root.value).toBe(3); + expect(tree.root.height).toBe(2); + + tree.insert(6); + + expect(tree.toString()).toBe('2,3,4,5,6'); + expect(tree.root.value).toBe(3); + expect(tree.root.right.value).toBe(5); + expect(tree.root.height).toBe(2); + }); + + it('should do complex right-right rotation', () => { + const tree = new AvlTree(); + + tree.insert(30); + tree.insert(20); + tree.insert(40); + tree.insert(50); + + expect(tree.root.value).toBe(30); + expect(tree.root.height).toBe(2); + expect(tree.toString()).toBe('20,30,40,50'); + + tree.insert(35); + expect(tree.root.value).toBe(30); + expect(tree.root.height).toBe(2); + expect(tree.toString()).toBe('20,30,35,40,50'); + + tree.insert(55); + expect(tree.root.value).toBe(40); + expect(tree.root.height).toBe(2); + expect(tree.toString()).toBe('20,30,35,40,50,55'); + }); + + it('should do left-right rotation', () => { + const tree = new AvlTree(); + + tree.insert(30); + tree.insert(20); + tree.insert(25); + + expect(tree.root.height).toBe(1); + expect(tree.root.value).toBe(25); + expect(tree.toString()).toBe('20,25,30'); + }); + + it('should do right-left rotation', () => { + const tree = new AvlTree(); + + tree.insert(30); + tree.insert(40); + tree.insert(35); + + expect(tree.root.height).toBe(1); + expect(tree.root.value).toBe(35); + expect(tree.toString()).toBe('30,35,40'); + }); + + it('should create balanced tree: case #1', () => { + const tree = new AvlTree(); + + tree.insert(1); + tree.insert(2); + tree.insert(3); + + expect(tree.root.value).toBe(2); + expect(tree.root.height).toBe(1); + expect(tree.toString()).toBe('1,2,3'); + + tree.insert(6); + + expect(tree.root.value).toBe(2); + expect(tree.root.height).toBe(2); + expect(tree.toString()).toBe('1,2,3,6'); + + tree.insert(15); + + expect(tree.root.value).toBe(2); + expect(tree.root.height).toBe(2); + expect(tree.toString()).toBe('1,2,3,6,15'); + + tree.insert(-2); + + expect(tree.root.value).toBe(2); + expect(tree.root.height).toBe(2); + expect(tree.toString()).toBe('-2,1,2,3,6,15'); + + tree.insert(-5); + + expect(tree.root.value).toBe(2); + expect(tree.root.height).toBe(2); + expect(tree.toString()).toBe('-5,-2,1,2,3,6,15'); + + tree.insert(-8); + + expect(tree.root.value).toBe(2); + expect(tree.root.height).toBe(3); + expect(tree.toString()).toBe('-8,-5,-2,1,2,3,6,15'); + }); +});