Build & traverse a parent-child relationship tree in dfs fashion

traverse_hierarchy_dfs.py

from collections import defaultdict

data = [(1, 'parent item', 0, 'chris co.'),
        (2, 'A. first child of main parent', 1, 'chris co.'),
        (3, 'B. second child of main parent', 1, 'chris co.'),
        (4, 'C. third child of main parent', 1, 'chris co.'),
        (5, 'D. only child of node A', 2, 'chris co.'),
        (6, 'E. first child of node B', 3, 'chris co.'),
        (7, 'F. second child of node B', 3, 'chris co.'),
        (8, 'G. only child of node C', 4, 'chris co.'),
        (9, 'H. first child of node F', 7, 'chris co.'),
        (10, 'I. 2nd child of node F', 7, 'chris co.'),
        (11, 'J. 3rd child of node F', 7, 'chris co.')]


def get_company_info(company_name: str):
    """
    Fake version of DB access.
    """
    links = [(x[0], x[2]) for x in data if x[3] == company_name]
    root = [x[0] for x in data if x[2] == 0][0]
    lookup = {x[0]: x[1] for x in data}

    return root, links, lookup


def build_relationship_tree(children):
    rel_tree = defaultdict(lambda: [])
    for item_id, item_parent in children:
        rel_tree[item_parent].append(item_id)
    return rel_tree


def stringify_child(lookup, tree, root_id, indent=0):
    s = [("\t" * indent) + lookup[root_id]]
    for child in tree[root_id]:
        s.extend(stringify_child(lookup, tree, child, indent+1))
    return s


root_id, links, lookup = get_company_info('chris co.')
tree = build_relationship_tree(links)
s_array = stringify_child(lookup, tree, root_id)
print(s_array)