From a0fa09a381c896821878cb7af74969123c70b2a5 Mon Sep 17 00:00:00 2001
From: Bertrand Benjamin <benjamin.bertrand@opytex.org>
Date: Wed, 31 Jan 2018 17:18:52 +0300
Subject: [PATCH] Balancing a AssocialTree

---
 mapytex/calculus/core/tree.py | 96 +++++++++++++++++++++++++++++++++++
 1 file changed, 96 insertions(+)

diff --git a/mapytex/calculus/core/tree.py b/mapytex/calculus/core/tree.py
index d9b0419..8a82cc8 100644
--- a/mapytex/calculus/core/tree.py
+++ b/mapytex/calculus/core/tree.py
@@ -966,6 +966,102 @@ class AssocialTree(Tree):
         except AttributeError:
             yield callback(self.right_value)
 
+    def balance(self):
+        """ Balance the tree
+
+        :example:
+        >>> t = AssocialTree.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 = AssocialTree.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
+
+
+        :returns: Balanced Tree (not AssocialTree)
+
+        """
+        leafs = self.get_leafs()
+        balanced_leafs = []
+        for l in leafs:
+            try:
+                balanced_leafs.append(l.balance())
+            except AttributeError:
+                balanced_leafs.append(l)
+
+        t = Tree.from_list(self.node, balanced_leafs)
+        return t
+
 # -----------------------------
 # Reglages pour 'vim'
 # vim:set autoindent expandtab tabstop=4 shiftwidth=4: