implement an arbitrary tree python

Question

It seems no similar topic exists.

I am required to build a tree from a line of integers for future processing. This line contains 𝑛 integer numbers from −1 to 𝑛 − 1 — parents of nodes. If the 𝑖-th one of them (0 ≤ 𝑖 ≤ 𝑛 − 1) is −1, node 𝑖 is the root, otherwise it's 0-based index of the parent of 𝑖-th node. eg

 the input line of integers: 4 -1 4 1 1 then -> 0 1 2 3 4 then -> 1 is the root, 3 and 4 are children of node 1, 0 and 2 are children of node 4. 

I should be able to do the remaining job but not very clear about the processing of constructing this tree.

Answer 1

You could use dict where values would be list of children to represent the tree. This would make construction of the tree trivial since for every node you can just append the current index to parent children. Here's an example using defaultdict with list as default factory:

from collections import defaultdict

s = '4 -1 4 1 1'

tree = defaultdict(list)
for node, parent in enumerate(int(x) for x in s.split()):
    tree[parent].append(node)

# Remove "parent" of the root node
root = tree.pop(-1)[0]

def print_tree(root, tree, prefix=''):
    print(prefix + str(root))
    for child in tree[root]:
        print_tree(child, tree, prefix + '  ')

print_tree(root, tree)

Output:

1
  3
  4
    0
    2

Update : Here's an additional example of tree traversal which counts the number of nodes in the tree. For every node visited it returns 1 + sum of child nodes:

def count_nodes(root, tree):
    return 1 + sum(count_nodes(child, tree) for child in tree[root])