Find the width of tree at each level/height (non-binary tree)-CodePudding

Dear experienced friends, I am looking for an algorithm (Python) that outputs the width of a tree at each level. Here are the input and expected outputs.

(I have updated the problem with a more complex edge list. The original question with sorted edge list can be elegantly solved by @Samwise answer.)

Input (Edge List: source-->target)

[[11,1],[11,2],
[10,11],[10,22],[10,33],
[33,3],[33,4],[33,5],[33,6]]

The tree graph looks like this:

       10
    /   |  \
   11   22  33
   / \    / | \ \
  1   2  3  4  5 6

Expected Output (Width of each level/height)

[1,3,6] # according to the width of level 0,1,2

I have looked through the web. It seems this topic related to BFS and Level Order Traversal. However, most solutions are based on the binary tree. How can solve the problem when the tree is not binary (e.g. the above case)?

(I'm new to the algorithm, and any references would be really appreciated. Thank you!)

CodePudding user response：

Build a dictionary of the "level" of each node, and then count the number of nodes at each level:

>>> from collections import Counter
>>> def tree_width(edges):
...     levels = {}  # {node: level}
...     for [p, c] in edges:
...         levels[c] = levels.setdefault(p, 0)   1
...     widths = Counter(levels.values())  # {level: width}
...     return [widths[level] for level in sorted(widths)]
...
>>> tree_width([[0,1],[0,2],[0,3],
... [1,4],[1,5],
... [3,6],[3,7],[3,8],[3,9]])
[1, 3, 6]

CodePudding user response：

This might not be the most efficient, but it requires only two scans over the edge list, so it's optimal up to a constant factor. It places no requirement on the order of the edges in the edge list, but does insist that each edge be (source, dest). Also, doesn't check that the edge list describes a connected tree (or a tree at all; if the edge list is cyclic, the program will never terminate).

from collections import defauiltdict

# Turn the edge list into a (non-binary) tree, represented as a
# dictionary whose keys are the source nodes with the list of children
# as its value.
def edge_list_to_tree(edges):
    '''Given a list of (source, dest) pairs, constructs a tree.
       Returns a tuple (tree, root) where root is the root node
       and tree is a dict which maps each node to a list of its children.
       (Leaves are not present as keys in the dictionary.)
    ''' 
    tree = defaultdict(list)
    sources = set()  # nodes used as sources
    dests = set()    # nodes used as destinations
    for source, dest in edges:
        tree[source].append(dest)
        sources.add(source)
        dests.add(dest)
    roots = sources - dests      # Source nodes which are not destinations
    assert(len(roots) == 1)      # There is only one in a tree
    tree.default_factory = None  # Defang the defaultdict
    return tree, roots.pop()

# A simple breadth-first-search, keeping the count of nodes at each level.
def level_widths(tree, root):
    '''Does a BFS of tree starting at root counting nodes at each level.
       Returns a list of counts.
    '''
    widths = []         # Widths of the levels
    fringe = [root]     # List of nodes at current level
    while fringe:
        widths.append(len(fringe))
        kids = []       # List of nodes at next level
        for parent in fringe:
            if parent in tree:
                for kid in tree[parent]:
                    kids.append(kid)
        fringe = kids   # For next iteration, use this level's kids
    return widths

# Put the two pieces together.
def tree_width(edges):
    return level_widths(*edge_list_to_tree(edges))