Python 从字典构造多叉树

Python 从字典构造多叉树

---

题目要求

构造一个树节点类，每个节点都有一个int类型的值（val），和0个或多个子节点childs。
给定一个字典，其中字典的key为当前节点的val，value对应列表对应当前节点的childs的val值。

请将该字典的构造为多叉树结构，使用root表示数的根节点，并打印每个节点的值及其子节点的值。（这里确保OrderDict字典的第一个key为树的根节点）。

输入：

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

---

代码实现

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):

        # 深度优先遍历构造多叉树结构
        def dfs(node, input_dict):
            if node.val in input_dict:  # 当前节点还有子节点时，继续对每个child向下构造
                for val in input_dict[node.val]:
                    cur_node = TreeNode(val)
                    dfs(cur_node, input_dict)
                    node.add_childs(cur_node)
            else:  # 当前节点没有子节点，直接返回
                return

        # 获取字典的第一个key，并以此为根节点构造多叉树
        root = TreeNode(list(input_dict.keys())[0])
        dfs(root, input_dict)

        return root

    @classmethod
    def print_tree(cls, root: "TreeNode"):
        """按照字典输入形式打印输出"""
        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"):
        """按照树的层级打印输出"""
        node_with_level = defaultdict(list)
        node_with_level[0].append([root])
        cur_level = 0

        while node_with_level:
            # 打印当前层节点
            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 node.childs:
                        node_with_level[cur_level + 1].append(node.childs)
                    else:  # 没有子节点的话，使用[]占位
                        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)

运行结果：

# print_tree方法输出字典格式结果：
1 : [2 3 4]
2 : [5 6]
5 : []
6 : []
3 : [7]
7 : [8 9 10]
8 : []
9 : []
10 : []
4 : []

# print_tree_graph方法输出类似数层级结构格式的结果：
[1] 
[2, 3, 4] 
[5, 6] [7] [] 
[] [] [8, 9, 10] 
[] [] []

对于第二种输出格式，其关系为：

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