How to get leaf nodes of a tree using Python?

Question

Hi there I am new to OOP so have this in mind while you are reading this.

I have a simple Python tree implementation(see below code).

class TreeNode(object):
    def __init__(self, data):
        self.data = data
        self.children = []

    def add_child(self, obj):
        self.children.append(obj)

class Tree:
    def __init__(self):
        self.root = TreeNode('ROOT')

    def preorder_trav(self, node):
        if node is not None:
            print node.data
            if len(node.children) == 0:
                print "("+ node.data + ")"
                for n in node.children:
                    self.preorder_trav(n)

if __name__ == '__main__':
    tr = Tree()
    n1 = tr.root
    n2 = TreeNode("B")
    n3 = TreeNode("C")
    n4 = TreeNode("D")
    n5 = TreeNode("E")
    n6 = TreeNode("F")

    n1.add_child(n2)
    n1.add_child(n3)
    n2.add_child(n4)
    n2.add_child(n5)
    n3.add_child(n6)

    tr.preorder_trav(n1)

What I need now is to implement a method for getting Leaf Nodes back. By the term leaf node I mean a node that has no children.

I am wondering how to make a get_leaf_nodes() method.

Some solutions come to my mind are

1. Making a self.leaf_nodes = [] inside the __init__ method. By making this I know it will be seen only by this tree instance.
2. Making a class member leaf_nodes = [] above __init__ method. By making this I know all tree instances will be able to see leaf_nodes list.

The above solutions will cause me to create a leaf_nodes list inside my class so the get_leaf_nodes() method could use. What I am looking for is to only have a get_leaf_nodes() method that will do the computation on my tree and will return a list.

For example in C we would call malloc() and then we could return the pointer to the function that called the get_leaf_nodes().

Answer 1

In python you can use an internal function to collect the leaf nodes and then return the list of them.

def get_leaf_nodes(self):
    leafs = []
    def _get_leaf_nodes( node):
        if node is not None:
            if len(node.children) == 0:
                leafs.append(node)
            for n in node.children:
                _get_leaf_nodes(n)
    _get_leaf_nodes(self.root)
    return leafs

If you want a more clean OOP approach you can create an extra private method for the collection of leafs:

def get_leaf_nodes(self):
    leafs = []
    self._collect_leaf_nodes(self.root,leafs)
    return leafs

def _collect_leaf_nodes(self, node, leafs):
    if node is not None:
        if len(node.children) == 0:
            leafs.append(node)
        for n in node.children:
            self._collect_leaf_nodes(n, leafs)

This is the way I'd do it in Java.

Answer 2

This method should be enough to get the leaves reachable from any node, if you call it with the root of your tree, you will obtain all of the leaves of the tree:

def get_leaves(node):
    if not node.children:
        yield node

    for child in node.children:
        for leaf in get_leaves(child):
             yield leaf