How to convert a nested list to a tree representation in Python?

Question

I am trying to create a tree visualiser from nested list using python without any special libraries.

The problem I encountered is that the tree, while its nodes are stored correctly, prints out in the order of the input.

I have tried implementing a recursive solution to this problem, which works for inputs such as [42,['a','b']] - where 42 is the root node and 'a' and 'b' are its children. However I can also receive input like [['a','b'], 42], where 42 is still the root node. Code, example output are below.

Here is the solution I tried implementing.

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)

This given code works for inputs such as [42,['a,'b']], where the output is

42
  a
  b

However for inputs like [['a','b'],42], the output is

  a
  b
42

Answer 1

I played around with your code a bit and discovered that what actually happens is that the "root node" has data of None, and what looks like the node's data is actually a child with data set and no children.

What you may be intending is this:

    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)

This would work if you know your input is always going to be 2-element lists, with at most one sub-list indicating the children.

In the case of something like @trincot's example [['a', [9], 'b'],42], depending on whether or not this makes sense, all you may have to do is this instead to make sure the leaf nodes are always listed first in the children list:

    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)

Take your pick. Or maybe someone else will post another answer.

Answer 2

What if we resort the node data by type as we process it:

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

OUTPUT

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