limitCracker limitCracker - 2 months ago 153
Python Question

How to get leaf nodes of a tree using Python?

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

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.