Python Construct a multitree from a dictionary

Subject requirements

Construct a tree node class , Each node has a int Type value （val）, and 0 One or more child nodes childs.
Give a dictionary , Among them, the dictionary of key Is the current node's val,value The corresponding list corresponds to the... Of the current node childs Of val value .

Please construct the dictionary as a multi tree structure , Use root Represents the root node of the number , And print the values of each node and its child nodes .（ Make sure that OrderDict The first one in the dictionary key Is the root node of the tree ）.

Code implementation

from collections import defaultdict


class TreeNode(object):
    def __init__(self, val=None) -> None:
        self.val = val
        self.childs = []

    def add_childs(self, child: "TreeNode"):
        self.childs.append(child)

    def __repr__(self) -> str:
        return str(self.val)

    @classmethod
    def build_tree(cls, input_dict: dict):

        #  Depth first traversal constructs a multi fork tree structure 
        def dfs(node, input_dict):
            if node.val in input_dict:  #  When the current node has child nodes , Continue to child Construct downward 
                for val in input_dict[node.val]:
                    cur_node = TreeNode(val)
                    dfs(cur_node, input_dict)
                    node.add_childs(cur_node)
            else:  #  The current node has no children , Go straight back to 
                return

        #  Get the first... Of the dictionary key, And take this as the root node to construct a multitree 
        root = TreeNode(list(input_dict.keys())[0])
        dfs(root, input_dict)

        return root

    @classmethod
    def print_tree(cls, root: "TreeNode"):
        """ Print out as dictionary input """
        if root:
            if root.childs:
                print(root.val, ": ", end="")
                print('[%s]' % (' '.join([str(_.val) for _ in root.childs])))
                for node in root.childs:
                    cls.print_tree(node)
            else:
                print(root.val, ": []")

    @classmethod
    def print_tree_graph(cls, root: "TreeNode"):
        """ Print the output according to the tree level """
        node_with_level = defaultdict(list)
        node_with_level[0].append([root])
        cur_level = 0

        while node_with_level:
            #  Print the current layer node 
            cur_nodes_lists = node_with_level.pop(cur_level)
            for nodes_list in cur_nodes_lists:
                print(nodes_list, end=" ")

                for node in nodes_list:  #  If there are children , Add it to the next layer 
                    if node.childs:
                        node_with_level[cur_level + 1].append(node.childs)
                    else:  #  If there are no child nodes , Use [] placeholder 
                        node_with_level[cur_level + 1].append([])
            cur_level += 1
            print()


input_dict = {
    
    1: [2, 3, 4],
    2: [5, 6],
    3: [7],
    7: [8, 9, 10],
}

tree = TreeNode.build_tree(input_dict)
TreeNode.print_tree(tree)
TreeNode.print_tree_graph(tree)

Running results ：

# print_tree Method to output dictionary format results ：
1 : [2 3 4]
2 : [5 6]
5 : []
6 : []
3 : [7]
7 : [8 9 10]
8 : []
9 : []
10 : []
4 : []

# print_tree_graph Method to output results in a number hierarchy format ：
[1] 
[2, 3, 4] 
[5, 6] [7] [] 
[] [] [8, 9, 10] 
[] [] []

For the second output format , The relationship is ：

           +---+
           | 1 |
           +-+-+
             |
             |
       +---+-v-+---+
  +----+ 2 | 3 | 4 +--+
  |    +---+-+-+---+  |
  |          |        |
  |          |        |
+-v-+---+   +v--+   +-v+
| 5 | 6 |   | 7 |   |  |
+-+-+--++   +-+-+   +--+
  |    |      |
  |    |      |
+-v+  +v-+  +-v-+---+---+
|  |  |  |  | 8 | 9 | 10|
+--+  +--+  ++--+-+-+--++
             |    |    |
             |    |    |
            +v-+ +v-+ +v-+
            |  | |  | |  |
            +--+ +--+ +--+