fix bug in the avl tree: root was not getting updated

Author: amejiarosario

src/data-structures/maps/map.spec.js

@@ -60,12 +60,11 @@ mapImplementations.forEach((MapImplementation) => {
       });
     });
 
-    fdescribe('#keys', () => {
+    describe('#keys', () => {
       beforeEach(() => {
         map.set(1, 2);
         map.set(2, 'dos');
         map.set(3, 3);
-        console.log(map.tree.toArray());
       });
 
       it('should return all keys', () => {

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

@@ -11,7 +11,7 @@ class AvlTree extends BinarySearchTree {
    */
   add(value) {
     const node = super.add(value);
-    balanceUptream(node);
+    this.root = balanceUptream(node);
     return node;
   }
 
@@ -23,7 +23,7 @@ class AvlTree extends BinarySearchTree {
     const node = super.find(value);
     if (node) {
       const found = super.remove(value);
-      balanceUptream(node.parent);
+      this.root = balanceUptream(node.parent);
       return found;
     }
 

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

@@ -29,6 +29,9 @@ describe('AvlTree', () => {
       expect(n.toValues()).toMatchObject({
         value: 2, left: 1, right: 3, parent: null,
       });
+
+      expect(tree.toArray()).toEqual([2, 1, 3, null, null, null, null]);
+      expect(tree.root).toBe(n);
     });
 
     it('should rebalance desc inserted nodes (RR Rotation)', () => {
@@ -39,6 +42,7 @@ describe('AvlTree', () => {
       expect(n.toValues()).toMatchObject({
         value: 2, left: 1, right: 3, parent: null,
       });
+      expect(tree.root).toBe(n);
     });
 
     it('should rebalance with LR Rotation', () => {
@@ -49,6 +53,7 @@ describe('AvlTree', () => {
       expect(n.toValues()).toMatchObject({
         value: 2, left: 1, right: 3, parent: null,
       });
+      expect(tree.root).toBe(n);
     });
 
 
@@ -60,6 +65,7 @@ describe('AvlTree', () => {
       expect(n.toValues()).toMatchObject({
         value: 2, left: 1, right: 3, parent: null,
       });
+      expect(tree.root).toBe(n);
     });
 
     it('should not balance', () => {
@@ -90,6 +96,9 @@ describe('AvlTree', () => {
       expect(n.toValues()).toMatchObject({
         value: 2, left: 1, right: 3, parent: null,
       });
+
+      expect(tree.root).toBe(n);
+      expect(tree.toArray()).toEqual([2, 1, 3, null, null, null, null]);
     });
 
     it('should rebalance desc inserted nodes (RR Rotation)', () => {
@@ -103,6 +112,7 @@ describe('AvlTree', () => {
       expect(n.toValues()).toMatchObject({
         value: 2, left: 1, right: 3, parent: null,
       });
+      expect(tree.root).toBe(n);
     });
 
     it('should rebalance with LR Rotation', () => {
@@ -116,6 +126,7 @@ describe('AvlTree', () => {
       expect(n.toValues()).toMatchObject({
         value: 2, left: 1, right: 3, parent: null,
       });
+      expect(tree.root).toBe(n);
     });
 
 
@@ -130,6 +141,7 @@ describe('AvlTree', () => {
       expect(n.toValues()).toMatchObject({
         value: 2, left: 1, right: 3, parent: null,
       });
+      expect(tree.root).toBe(n);
     });
   });
 });

src/data-structures/trees/binary-search-tree.spec.js

@@ -60,6 +60,7 @@ describe('Binary Search Tree', () => {
         const n3 = bst.add(3);
 
         expect(n3.parent.value).toBe(2);
+        expect(bst.toArray()).toEqual([1, null, 2, null, 3, null, null]);
       });
 
       it('should deal with duplicates', () => {

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

@@ -140,18 +140,19 @@ function balance(node) {
   if (node.balanceFactor > 1) {
     // left subtree is higher than right subtree
     if (node.left.balanceFactor > 0) {
-      rightRotation(node);
+      return rightRotation(node);
     } else if (node.left.balanceFactor < 0) {
-      leftRightRotation(node);
+      return leftRightRotation(node);
     }
   } else if (node.balanceFactor < -1) {
     // right subtree is higher than left subtree
     if (node.right.balanceFactor < 0) {
-      leftRotation(node);
+      return leftRotation(node);
     } else if (node.right.balanceFactor > 0) {
-      rightLeftRotation(node);
+      return rightLeftRotation(node);
     }
   }
+  return node;
 }
 
 /**
@@ -161,10 +162,12 @@ function balance(node) {
  */
 function balanceUptream(node) {
   let current = node;
+  let newParent;
   while (current) {
-    balance(current);
+    newParent = balance(current);
     current = current.parent;
   }
+  return newParent;
 }
 
 module.exports = {