Flatten a dictionary into a list containing the indices of parents in their generation

user_8881449 0
Viewed 53

I have an input dictionary that describes a tree structure, eg:

{'title': 'Root', 'children': [
    {'title': 'Child 1','children': [
        {'title': 'Grandchild 11', 'children': [
            {'title': 'Great Grandchild 111', 'children': []}
        ]}
    ]},
    {'title': 'Child 2', 'children': [
        {'title': 'Grandchild 21', 'children': []}
    ]},
    {'title': 'Child 3', 'children': [
        {'title': 'Grandchild 31', 'children': []}
    ]}
]}

I am trying to write a python function that flattens this structure into a list. The list should contain one int per object in this dictionary, where the integer represents the index of that object's parent in the parent's own 'generation'. The traversal of the tree should be level-order.

In the case of the example above, the expected output would be

[0, 0, 0, 0, 0, 1, 2, 0]

... as the parent indices corresponding to the objects titled:

['Root', 'Child 1', 'Child 2', 'Child 3', 'Grandchild 11', 'Grandchild 21', 'Grandchild 31', 'Great Grandchild 111']

This feels like a Lisp problem to me and Python recursion has me hung up.

tree recursion dictionary python
1 Answers

I'll assume you have this input (it seems you forgot to close the children list of node "Child 2"):

d = {'title': 'Root', 'children': [
    {'title': 'Child 1','children': [
        {'title': 'Grandchild 11', 'children': [
            {'title': 'Great Grandchild 111', 'children': []}
        ]}
    ]},
    {'title': 'Child 2', 'children': [
        {'title': 'Grandchild 21', 'children': []}
    ]},
    {'title': 'Child 3', 'children': [
        {'title': 'Grandchild 31', 'children': []}
    ]}
]}

Then perform a breadth-first traversal, keeping track of the index that each node has in its parent's children list:

def collect(d):
    q = [(0, d)]
    yield 0
    while q:
        front = []
        for idx, node in q:
            yield from [idx] * len(node["children"])
            front.extend(enumerate(node["children"]))
        q = front

Example use:

print(list(collect(d)))

Output:

[0, 0, 0, 0, 0, 1, 2, 0]
user_54598390
Related