tree over

Author: EINDEX

data_structure/tree/tree.py

@@ -17,6 +17,8 @@ def __init__(self, value, children: list = None):
 
         if not children:
             self.children = []
+        elif not self.value:
+            self.children = []
         else:
             self.children = children
 
@@ -43,17 +45,16 @@ def preorder_traversal_while(self):
         res = []
         if not self.root:
             return res
+        stack = [self.root]
 
-        queue = Queue()
-        queue.put(self.root)
-
-        while queue.qsize():
-            node = queue.get()
+        while len(stack):
+            node = stack.pop()
             if not node.value:
                 continue
+            for sub_node in node.children[::-1]:
+                stack.append(sub_node)
             res.append(node.value)
-            for sub_node in node.children:
-                queue.put(sub_node)
+
         return res
 
     def postorder_traversal_while(self):
@@ -115,11 +116,118 @@ def _inner(root):
 
         return _inner(self.root)
 
+    def layer_while(self):
+        """
+        树的层序遍历
+        :return: list of tree node value
+        """
+        res = []
+        if not self.root:
+            return res
+
+        queue = Queue()
+        queue.put(self.root)
+
+        while queue.qsize():
+            node = queue.get()
+            if not node.value:
+                continue
+            res.append(node.value)
+            for sub_node in node.children:
+                queue.put(sub_node)
+        return res
+
+    def layer_recursion(self):
+        # todo
+        pass
+
+    def depth_recursion(self):
+
+        def _inner(root, depth=1):
+            if not root.children:
+                return depth
+            return max([_inner(sub_node, depth+1) for sub_node in root.children])
+
+        return _inner(self.root)
+
+    def node_count(self):
+        def _inner(root):
+            if not root.children:
+                return 1
+            return 1 + sum([_inner(sub_node) for sub_node in root.children])
+        return _inner(self.root)
+
+    def leaf_count(self):
+        def _inner(root):
+            if not root.children:
+                return 1
+            return sum([_inner(sub_node) for sub_node in root.children])
+        return _inner(self.root)
+
+    def lowest_ancestor_node(self, node1, node2):
+        stack = [self.root]
+        stack1 = None
+        stack2 = None
+
+        while len(stack) and not (stack1 and stack2):
+            node = stack.pop()
+
+            if node is node1:
+                stack1 = stack[:]
+            if node is node2:
+                stack2 = stack[:]
+            if not node.value:
+                continue
+            for sub_node in node.children:
+                stack.append(sub_node)
+
+        res = self.root
+        for i in range(len(stack1)):
+            if stack1[i] == stack2[i]:
+                res = stack1[i]
+            else:
+                return res.value
+
+        return res.value
+
+    def two_node_distence(self, node1, node2):
+        stack = [self.root]
+        stack1 = None
+        stack2 = None
+
+        while len(stack) and not (stack1 and stack2):
+            node = stack.pop()
+
+            if node is node1:
+                stack1 = stack[:]
+            if node is node2:
+                stack2 = stack[:]
+            if not node.value:
+                continue
+            for sub_node in node.children:
+                stack.append(sub_node)
+
+        res = self.root
+        for i in range(len(stack1)):
+            if stack1[i] == stack2[i]:
+                res = stack1[i]
+            else:
+                return len(stack1) + len(stack2) - 2*i
+
+        return len(stack1) + len(stack2) - 2*i
+
 
 if __name__ == '__main__':
+    node1, node2 = TreeNode(6), TreeNode(8)
     t = Tree(
-        TreeNode(1, [TreeNode(2, [TreeNode(4), TreeNode(5), TreeNode(6)]), TreeNode(3, [TreeNode(7), TreeNode(8)])]))
+        TreeNode(1, [TreeNode(2, [TreeNode(4), TreeNode(5), node1]), TreeNode(3, [TreeNode(7), node2])]))
     print(t.preorder_traversal_while())
     print(t.postorder_traversal_while())
     print(t.preorder_traversal_recursion())
     print(t.postorder_traversal_recursion())
+    print(t.layer_while())
+    print(t.depth_recursion())
+    print(t.node_count())
+    print(t.leaf_count())
+    print(t.lowest_ancestor_node(node1, node2))
+    print(t.two_node_distence(node1, node2))