Counting internal (non-leaf) nodes in a tree starting from any node

Viewed 499

Here is a tree I have. I want to calculate how many internal nodes (non-leaf nodes) starting from any node in the tree.

example tree

The function that I have now can only count all that is below A which gives 10. But I only want those before the terminal node.

  • eg. for A, it should be 4 (B, C, E, H)
  • eg. for E, it should be 1 (H)

How can I modify my function to do that? Note that my data is very large so a recursive function would cause a stack overflow.

  name employee
0   A   B
1   A   C
2   B   D
3   C   E
4   C   F
5   E   H
6   E   I
7   H   T
8   H   U
9   H   V

#collapse to dictionary
dict_a = {k: g["employee"].tolist() for k,g in em.groupby("name")}
dict_a

{'A': ['B', 'C'],
 'B': ['D'],
 'C': ['E', 'F'],
 'E': ['H', 'I'],
 'H': ['T', 'U', 'V']}
    
# recursive function to count lengths:
def total(k,connections):
    if k not in connections:
        return 0
    
    # number of direct connections plus their connections:
    return len(connections[k]) + sum(total(child_k, connections) for child_k in connections[k])
2 Answers

Getting interior nodes means running your traversal and only yielding nodes with at least one child.

The caller has the flexibility to turn the generator into a list and take the length if desired, or simply iterate since the function is general (a count-specific function seems unnecessarily narrow, but if you need it, it's a simple wrapper on the below function).

def interior_nodes_from(tree, node):
    if node in tree:
        for child in tree[node]:
            if child in tree:
                yield child
                yield from interior_nodes_from(tree, child)

if __name__ == '__main__':
    tree = {'A': ['B', 'C'],
            'B': ['D'],
            'C': ['E', 'F'],
            'E': ['H', 'I'],
            'H': ['T', 'U', 'V']}
    
    for target in 'ABCDEFGH':
        interior_nodes = list(interior_nodes_from(tree, target))
        print(f'{target} => {len(interior_nodes)} children, {interior_nodes}')

Output:

A => 4 children, ['B', 'C', 'E', 'H']
B => 0 children, []
C => 2 children, ['E', 'H']
D => 0 children, []
E => 1 children, ['H']
F => 0 children, []
G => 0 children, []
H => 0 children, []

If your tree depth exceeds the call stack size, you can use an explicit stack instead of recursion:

def interior_nodes_from(tree, node):
    if node in tree:
        stack = [node]
        
        while stack:
            for child in tree[stack.pop()]:
                if child in tree:
                    yield child
                    stack.append(child)

One solution I would suggest in recursive function is :

  1. Check if the parameter k exists in connection, else return 1.

  2. Store the sum of each total(ck) in a variable inside the function, where ck is key of every child ( for loop ).

  3. Check if the sum == dict_a[k], if this is the case then that node is what you want, add it to the terminal node list.

    #Terminal Nodes termn = []

    #recursive function to count lengths: def total(k): if k not in dict_a: return 1

     # Number of child nodes.
     nc = 0;
     nc += total(child_k) for child_k in dict_a[k])
     if nc == len(dict_a[k]):
         termn.append(k)
    
Related