construct a python rich tree from list of descendants

Question

I am trying to construct a python rich.Tree from a list of descendants, d.

d = {0: [1, 2], 2: [3, 4, 5], 4: [6, 7]}

I expect to get something like this:

# manually constructed string
expected = \
"""
0
├── 1
└── 2
    └── 3
    └── 4
        └── 6
        └── 7
    └── 5
"""

But I am getting confused as to how to proceed with the construction: the following code is not correct.

from rich.tree import Tree
from rich import print as rprint

tree = Tree("0")
tree.add("1")
tree.add("2").add("3")
tree.add("4").add("6").add("7")
tree.add("5")
rprint(tree)

0
├── 1
├── 2
│   └── 3
├── 4
│   └── 6
│       └── 7
└── 5

Any suggestions will be appreciated, thanks!

Answer 1

You could use a recursive function to convert the dictionary representation of your tree (i.e. the adjacency list) to the nested tree object:

def make_tree(adj, node=0):
    tree = Tree(str(node))
    for child in adj.get(node, []):
        tree.add(make_tree(adj, child))
    return tree

Call like this:

d = {0: [1, 2], 2: [3, 4, 5], 4: [6, 7]}
tree = make_tree(d)
rprint(tree)