Feat(Tree): Write balance Tree method

mapytex/calculus/core/tree.py

@@ -611,14 +611,14 @@ class Tree(object):
             return self.right_value
 
     @property
-    def deep_branch(self):
+    def long_branch(self):
-        """ Get the deep branch of the tree, right if same depth
+        """ Get the long branch of the tree, right if same depth
         
         :return: A tree or a leaf value
 
         :example:
         >>> t = Tree.from_str("2+3*4")
-        >>> print(t.deep_branch)
+        >>> print(t.long_branch)
         *
          > 3
          > 4
@@ -631,12 +631,12 @@ class Tree(object):
          > *
          | > 4
          | > 5
-        >>> print(t.deep_branch)
+        >>> print(t.long_branch)
         *
          > 4
          > 5
         >>> t = Tree.from_str("2*3+4")
-        >>> print(t.deep_branch)
+        >>> print(t.long_branch)
         *
          > 2
          > 3
@@ -655,6 +655,90 @@ class Tree(object):
         else:
             return self.left_value
 
+    def balance(self):
+        """ Recursively balance the tree without permutting different nodes
+
+        :return: balanced tree
+
+        :example:
+        >>> t = Tree.from_str("1+2+3+4+5+6+7+8+9")
+        >>> print(t)
+        +
+         > +
+         | > +
+         | | > +
+         | | | > +
+         | | | | > +
+         | | | | | > +
+         | | | | | | > +
+         | | | | | | | > 1
+         | | | | | | | > 2
+         | | | | | | > 3
+         | | | | | > 4
+         | | | | > 5
+         | | | > 6
+         | | > 7
+         | > 8
+         > 9
+        >>> bal_t = t.balance()
+        >>> print(bal_t)
+        +
+         > +
+         | > +
+         | | > +
+         | | | > 1
+         | | | > 2
+         | | > 3
+         | > +
+         | | > 4
+         | | > 5
+         > +
+         | > 6
+         | > +
+         | | > 7
+         | | > +
+         | | | > 8
+         | | | > 9
+        >>> t = Tree.from_str("1+2+3+4+5*6*7*8*9")
+        >>> bal_t = t.balance()
+        >>> print(bal_t)
+        """
+        try:
+            l_depth = self.left_value.depth()
+        except AttributeError:
+            l_depth = 1
+        try:
+            r_depth = self.right_value.depth()
+        except AttributeError:
+            r_depth = 1
+
+        if l_depth > r_depth+1 and\
+                self.node == self.left_value.node:
+            new_left = self.left_value.long_branch
+            new_right = Tree(self.node,
+                             self.left_value.short_branch,
+                             self.right_value)
+            return Tree(self.node, new_left, new_right).balance()
+
+        if r_depth > l_depth+1 and\
+                self.node == self.right_value.node:
+            new_right = self.right_value.long_branch
+            new_left = Tree(self.node,
+                             self.left_value,
+                             self.right_value.short_branch)
+            return Tree(self.node, new_left, new_right).balance()
+        
+        try:
+            left_v = self.left_value.balance()
+        except AttributeError:
+            left_v = self.left_value
+        try:
+            right_v = self.right_value.balance()
+        except AttributeError:
+            right_v = self.right_value
+
+        return Tree(self.node, left_v, right_v)
+
 
 class MutableTree(Tree):
 
@@ -945,6 +1029,7 @@ class MutableTree(Tree):
 
         self.right_value = nright_value
 
+
 class LeafTree(Tree):
 
     """ A LeafTree is a tree which act as if it is a leaf.
@@ -961,6 +1046,7 @@ class LeafTree(Tree):
     def apply_on_last_level(self, *args):
         raise AttributeError("Can't use apply_on_last_level on a LeafTree")
 
+
 class AssocialTree(Tree):
 
     """ Tree which concider every subtree with a node different from itself