Tree utils

Author: suyash248

Tree/__init__.py

@@ -0,0 +1,25 @@
+from typing import TypeVar
+
+T = TypeVar('T')
+
+
+class TreeNode(object):
+    def __init__(self, data: T, left=None, right=None):
+        self.data = data
+        self.left = left
+        self.right = right
+
+    def __str__(self):
+        return "{left} <-- {data} --> {right}".format(
+            left="NULL" if self.left is None else self.left.data,
+            right="NULL" if self.right is None else self.right.data,
+            data=self.data
+        )
+
+
+def print_tree(root: TreeNode):
+    if root is None:
+        return
+    print(root)
+    print_tree(root.left)
+    print_tree(root.right)

test2.py

@@ -1,6 +1,8 @@
-from test import Node, print_tree
 from collections import deque
 
+from Tree import TreeNode, print_tree
+
+
 def deserialize(data):
     """Decodes your encoded data to tree.
 
@@ -12,18 +14,18 @@ def deserialize(data):
 
     data = data.split(",")
     q = deque()
-    root = Node(data[0])
+    root = TreeNode(data[0])
     q.append(root)
     i = 1
 
     while i < len(data):
         parent = q.popleft()
         if data[i] != '#':
-            parent.left = Node(data[i])
+            parent.left = TreeNode(data[i])
             q.append(parent.left)
         i += 1
         if data[i] != '#':
-            parent.right = Node(data[i])
+            parent.right = TreeNode(data[i])
             q.append(parent.right)
         i += 1
     return root