fix(avl-tree): balance was not working properly

Fixes: amejiarosario#33

Author: amejiarosario

src/data-structures/trees/avl-tree.js

@@ -17,23 +17,21 @@ const {
  * - LR rotations: double rotation left-right
  * - RL rotations: double rotation right-left
  *
- * @param {TreeNode} node
+ * @param {BinaryTreeNode} node
  */
 function balance(node) {
   if (node.balanceFactor > 1) {
     // left subtree is higher than right subtree
-    if (node.left.balanceFactor > 0) {
-      return rightRotation(node);
-    } if (node.left.balanceFactor < 0) {
+    if (node.left.balanceFactor < 0) {
       return leftRightRotation(node);
     }
-  } else if (node.balanceFactor < -1) {
+    return rightRotation(node);
+  } if (node.balanceFactor < -1) {
     // right subtree is higher than left subtree
-    if (node.right.balanceFactor < 0) {
-      return leftRotation(node);
-    } if (node.right.balanceFactor > 0) {
+    if (node.right.balanceFactor > 0) {
       return rightLeftRotation(node);
     }
+    return leftRotation(node);
   }
   return node;
 }
@@ -43,7 +41,7 @@ function balance(node) {
 /**
  * Bubbles up balancing nodes a their parents
  *
- * @param {TreeNode} node
+ * @param {BinaryTreeNode} node
  */
 function balanceUpstream(node) {
   let current = node;

src/data-structures/trees/avl-tree.spec.js

@@ -165,4 +165,65 @@ describe('AvlTree', () => {
         null, null, null, null, null, null]);
     });
   });
+
+  describe('balancing to the left', () => {
+    let n32;
+    beforeEach(() => {
+      n32 = tree.add(32);
+      tree.add(8);
+      tree.add(64);
+      tree.add(4);
+      tree.add(16);
+      tree.add(48);
+      tree.add(128);
+      tree.add(2);
+      tree.add(6);
+      tree.add(10);
+      tree.add(20);
+    });
+
+    it('should have all nodes', () => {
+      expect(tree.toArray()).toEqual([32, 8, 64, 4, 16, 48, 128, 2, 6, 10, 20,
+        null, null, null, null, null, null, null, null, null, null, null, null]);
+    });
+
+    it('should rebalance when removing', () => {
+      tree.remove(64);
+      expect(tree.toArray()).toEqual([32, 8, 128, 4, 16, 48, null, 2, 6, 10, 20,
+        null, null, null, null, null, null, null, null, null, null]);
+      expect(n32.balanceFactor).toBe(1);
+      expect(n32.right.balanceFactor).toBe(1);
+      expect(n32.left.balanceFactor).toBe(0);
+
+      tree.remove(48);
+      expect(tree.toArray()).toEqual([8, 4, 32, 2, 6, 16, 128, null, null, null, null, 10, 20,
+        null, null, null, null, null, null]);
+    });
+  });
+
+  describe('balancing to the right', () => {
+    beforeEach(() => {
+      tree.add(8);
+      tree.add(4);
+      tree.add(32);
+      tree.add(2);
+      tree.add(16);
+      tree.add(64);
+      tree.add(10);
+      tree.add(20);
+      tree.add(60);
+      tree.add(70);
+    });
+
+    it('should build the tree', () => {
+      expect(tree.toArray()).toEqual([8, 4, 32, 2, null, 16, 64, null, null, 10, 20, 60, 70,
+        null, null, null, null, null, null, null, null]);
+    });
+
+    it('should rebalance right side', () => {
+      tree.remove(2);
+      expect(tree.toArray()).toEqual([32, 8, 64, 4, 16, 60, 70, null, null, 10, 20,
+        null, null, null, null, null, null, null, null]);
+    });
+  });
 });