如何在Python中将嵌套列表转换为树形表示？

Question

我正在尝试用 Python 创建一个树形可视化工具，使用的是嵌套列表，并且不想用任何特别的库。

我遇到的问题是，虽然树的节点存储得很正确，但打印出来的顺序却是按照输入的顺序。

我尝试实现一个递归的解决方案，这个方案对像 [42,['a','b']] 这样的输入有效——在这里 42 是根节点，而 'a' 和 'b' 是它的子节点。不过，我也可能会收到像 [['a','b'], 42] 这样的输入，在这种情况下 42 仍然是根节点。下面是代码和示例输出。

这是我尝试实现的解决方案。

class TreeNode:
    def __init__(self, data):
        self.data = data
        self.children = []

    def add_child(self, child_node):
        self.children.append(child_node)

def assign_tree_nodes(input_data):
    if isinstance(input_data, list):
        node = TreeNode(None)
        for item in input_data:
            child_node = assign_tree_nodes(item)
            node.add_child(child_node)
        return node
    else:
        return TreeNode(input_data)

def print_tree(node, indent=0):
    if node.data is not None:
        print("  " * indent + str(node.data))
    if node.children:
        for child in node.children:
            print_tree(child, indent + 1)

# Prompt the user for input
input_data = eval(input("INPUT:\n"))

# Assign tree nodes
root_node = assign_tree_nodes(input_data)


# Print out the tree
for child in root_node.children:
    print_tree(child)

这个代码对像 [42,['a','b']] 这样的输入有效，输出结果是

42
  a
  b

但是对于像 [['a','b'],42] 这样的输入，输出结果是

  a
  b
42

Answer 1

如果我们在处理节点数据时，按照类型重新排序这些数据，会发生什么呢：

def assign_tree_nodes(input_data):
    if isinstance(input_data, list):
        data, children = sorted(input_data, key=lambda x: [int, list].index(type(x)))

        node = TreeNode(data)

        for child in children:
            node.add_child(assign_tree_nodes(child))
    else:
        node = TreeNode(input_data)

    return node

def print_tree(node, indent=0):
    print("  " * indent, node.data, sep='')

    for child in node.children:
        print_tree(child, indent + 1)

example = [6, [[[[1, [2, 3]], [[-43, 44], 42]], 4], 5]]

print_tree(assign_tree_nodes(example))

输出结果

% python3 test.py
6
  4
    1
      2
      3
    42
      -43
      44
  5
% 

Answer 2

我稍微玩了一下你的代码，发现其实发生的事情是“根节点”的数据是 None，而看起来像是节点数据的其实是一个有数据但没有子节点的子节点。

你可能想要的是这个：

    def add_child(self, child_node):
        if child_node.data is not None and child_node.children == []:
            # whoops, it's a leaf
            self.data = child_node.data
        else:
            self.children.append(child_node)

如果你知道输入总是会是包含两个元素的列表，并且最多只有一个子列表表示子节点的话，这样是可以工作的。

比如在 @trincot 的例子 [['a', [9], 'b'],42] 中，具体怎么处理取决于这是否有意义。你可能只需要这样做，确保叶子节点总是排在子节点列表的最前面：

    def add_child(self, child_node):
        if child_node.data:
            # it's a leaf so it comes first
            self.children.insert(0, child_node)
        else:
            # it's a branch node
            self.children.append(child_node)

你可以选择其中一个。或者也许其他人会给出不同的答案。