Search code examples
pythondictionarytreebreadth-first-search

Python: Number of nodes per level in dictionary with breadth-first search


Assuming the input

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': []}
        ]}
    ]}

I'm trying to write a python function that accepts d and returns a list of integers, where each integer represents the number of nodes per level of the dictionary as found in a breadth-first search.

In the case of the above example, I'd expect the output: [1, 3, 3, 1]


Solution

  • This can indeed be done with a breadth-first traversal:

    def levelwidths(d):
        level = [d]
        while level:
            yield len(level)
            level = [child for node in level for child in node["children"]]
    

    Example run:

    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': []}
        ]}
    ]}
    
    print(list(levelwidths(d)))
    

    Output:

    [1, 3, 3, 1]