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

Answer 3

One of the good recursive approach to find leaves.

def leaf_count(self):

        if(self == None):
            return 0

        if(self.left == None or self.right == None):
            return 1

        return self.left.leaf_count() + self.right.leaf_count()

Answer 4

I think previous contributors have responded to this question correctly but no one really showed how to setup the nodes and add those values to the tree and then running it to demonstrate it actually works all together. That's why I'm answering this question:

class Node: 
    def __init__(self, val): 
        self.val = val
        self.left = None
        self.right = None


class Pattern():
    def getLeafs(self,root):
        if not root: 
            return []
        if not root.left and not root.right: 
            return [root.val]
        leaves = self.getLeafs(root.left) + self.getLeafs(root.right)
        return leaves 

    def similar_leaf(self, root1, root2):
        return self.getLeafs(root1) == self.getLeafs(root2)


#simple test 
root= Node(3)
root.left = Node(5)
root.right= Node(1)

root.left.left =  Node(6)
root.left.right =  Node(2)
root.left.right.left =Node(7)
root.left.right.right =Node(4)

root.right.left =  Node(9)
root.right.right =  Node(8)

pattern =Pattern()
print(pattern.similar_leaf(root,root))

Answer 5

Approach below works fine for the sub trees which doesn't have left nodes or a right node.

Node

class Node:
    
    def __init__(self, data, left=None, right=None):
        self.data = data
        self.left = left
        self.right = right

As a Utility function

def get_leaf_nodes(root):
    leaves = []
    
    # If no child nodes
    if not root.left and not root.right:
        return [root.data]
    
    # If no any left child
    if not root.left:
        leaves = get_leaf_nodes(root.right)
    
    # If no any right child
    if not root.right:
        leaves = get_leaf_nodes(root.left)
    
    # If has left as well left child
    if root.left and root.right:
        leaves = get_leaf_nodes(root.left) + get_leaf_nodes(root.right)
    
    return leaves

As a class method:

class Node:
    
    def __init__(self, data, left=None, right=None):
        self.data = data
        self.left = left
        self.right = right

    def get_leaf_nodes(self):
        """Get the nodes which don't have any left of right sub nodes"""
        leaves = []
      
        # If no child nodes
        if not self.left and not self.right:
            return [self.data]
      
        # If no any left child
        if not self.left:
            leaves = self.right.get_leaf_nodes()
      
        # If no any right child
        if not self.right:
            leaves = self.left.get_leaf_nodes()
      
        # If has left as well left child
        if self.left and self.right:
            leaves = self.left.get_leaf_nodes() + self.right.get_leaf_nodes()
      
        return leaves

Test: Consider the binary tree as:

tree = Node(2,
            Node(1),
            Node(6,
                Node(3),
                Node(11,
                    Node(10,
                        Node(8,
                            None,
                            Node(9)
                            ),
                        None
                    )
                )
            )
    )

Then the leaves are:

# as a utility function
l = get_leaf_nodes(tree)

# or as an object
# l = tree.get_leaf_nodes()
print(l)

[1, 3, 9]