day18: trying to learn/play with binary trees

day18.py

@@ -73,8 +73,10 @@ class Snailfish:
 
 
 class BinarySnailfish:
-    def __init__(self, value: Union[int, List] = None, depth: int = 0):
+    def __init__(self, value: Union[int, List] = None, depth: int = 0, parent: 'BinarySnailfish' = None):
+        self.fours = []
         self.depth = depth
+        self.parent = parent
         if not isinstance(value, list):
             self.value = value
             self.left = None
@@ -83,41 +85,121 @@ class BinarySnailfish:
             self.value = None
             if isinstance(value[0], BinarySnailfish):
                 self.left = value[0]
-                self.left.incDepth()
+                self.left.setDepth(self.depth + 1)
+                self.left.parent = self
             else:
-                self.left = BinarySnailfish(value[0])
+                self.left = BinarySnailfish(value[0], parent=self, depth=self.depth + 1)
 
             if isinstance(value[1], BinarySnailfish):
                 self.right = value[1]
-                self.right.incDepth()
+                self.right.setDepth(self.depth + 1)
+                self.right.parent = self
             else:
-                self.right = BinarySnailfish(value[1])
+                self.right = BinarySnailfish(value[1], parent=self, depth=self.depth + 1)
 
-        while reduce(self):
-            pass
-
-    def incDepth(self):
-        self.depth += 1
-        if self.value:
+    def setDepth(self, depth: int):
+        self.depth = depth
+        if self.value is not None:
             return
         else:
-            self.left.incDepth()
-            self.right.incDepth()
+            if self.depth == 4:
+                self.getRoot().fours.append(self)
+            self.left.setDepth(depth + 1)
+            self.right.setDepth(depth + 1)
+
+    def getRoot(self) -> 'BinarySnailfish':
+        root = self
+        while root.parent:
+            root = root.parent
+
+        return root
+
+    def getLeftNumber(self) -> Union[None, 'BinarySnailfish']:
+        leftnum = None
+        for current in self.getRoot().traverse_inorder():
+            if current == self:
+                return leftnum
+            if current.value is not None:
+                leftnum = current
+
+    def getRightNumber(self) -> Union[None, 'BinarySnailfish']:
+        returnNext = 0
+        for current in self.getRoot().traverse_inorder():
+            if returnNext >= 3 and current.value is not None:
+                return current
+            if returnNext > 0:
+                returnNext += 1
+            if current == self:
+                returnNext += 1
+
+    def split(self):
+        self.left = BinarySnailfish(self.value // 2, depth=self.depth + 1, parent=self)
+        self.right = BinarySnailfish(int(ceil(self.value / 2)), depth=self.depth + 1, parent=self)
+        self.value = None
+
+    def explode(self):
+        left_num = self.getLeftNumber()
+        if left_num is not None and left_num.value is not None:
+            left_num.value += self.left.value
+
+        right_num = self.getRightNumber()
+        if right_num is not None and right_num.value is not None:
+            right_num.value += self.right.value
+
+        self.left = self.right = None
+        self.value = 0
+
+    def reduce(self):
+        print(self.fours)
+        if self.value is not None:
+            return
+        found = True
+        while found:
+            found = False
+            for n in self.fours:
+                self.fours.pop(0)
+                n.explode()
+                found = True
+                break
+
+            if found:
+                continue
+
+            for n in self.traverse_inorder():
+                if n.value is not None and n.value > 9:
+                    n.split()
+                    found = True
+                    break
+
+    def traverse_inorder(self):
+        yield self
+
+        if self.left:
+            for x in self.left.traverse_inorder():
+                yield x
+
+        if self.right:
+            for x in self.right.traverse_inorder():
+                yield x
 
     def __add__(self, other):
-        return BinarySnailfish([self, other])
+        sum_fish = BinarySnailfish([self, other])
+        sum_fish.reduce()
+        return sum_fish
+
+    def __str__(self):
+        if self.value is None:
+            return "[%s,%s]" % (str(self.left), str(self.right))
+        else:
+            return str(self.value)
 
     def getMagnitude(self):
-        if self.value:
+        if self.value is not None:
             return self.value
         else:
             return self.left.getMagnitude() * 3 + self.right.getMagnitude() * 2
 
 
-def reduce(node: BinarySnailfish):
-    pass
-
-
 class Day(AOCDay):
     test_solutions_p1 = [4140, 4417]
     test_solutions_p2 = [3993, 4796]