Python Language Tree exploration with recursion


Example

Say we have the following tree:

root
- A
  - AA
  - AB
- B
  - BA
  - BB
    - BBA

Now, if we wish to list all the names of the elements, we could do this with a simple for-loop. We assume there is a function get_name() to return a string of the name of a node, a function get_children() to return a list of all the sub-nodes of a given node in the tree, and a function get_root() to get the root node.

root = get_root(tree)
for node in get_children(root):
    print(get_name(node))
    for child in get_children(node):
        print(get_name(child))
        for grand_child in get_children(child):
            print(get_name(grand_child))
# prints: A, AA, AB, B, BA, BB, BBA

This works well and fast, but what if the sub-nodes, got sub-nodes of its own? And those sub-nodes might have more sub-nodes... What if you don't know beforehand how many there will be? A method to solve this is the use of recursion.

def list_tree_names(node):
    for child in get_children(node):
        print(get_name(child))
        list_tree_names(node=child)

list_tree_names(node=get_root(tree))
# prints: A, AA, AB, B, BA, BB, BBA

Perhaps you wish to not print, but return a flat list of all node names. This can be done by passing a rolling list as a parameter.

def list_tree_names(node, lst=[]):
    for child in get_children(node):
        lst.append(get_name(child))
        list_tree_names(node=child, lst=lst)
    return lst

list_tree_names(node=get_root(tree))
# returns ['A', 'AA', 'AB', 'B', 'BA', 'BB', 'BBA']