Add AVL Tree.

Author: trekhleb

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);
+  }
+}

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');
+  });
+});