Program to find leaf and non-leaf nodes of a binary tree in Python

In a binary tree, a leaf node has no children, while a non-leaf node has at least one child. This program counts both types of nodes using a recursive approach that traverses the entire tree.

So, if the input is like:

then the output will be (3, 2), as there are 3 leaf nodes and 2 non-leaf nodes.

Algorithm

To solve this, we will follow these steps ?

• If the node is null, return (0, 0)
• If the node has no children (leaf node), return (1, 0)
• Recursively count leaves and non-leaves in left and right subtrees
• Return the sum of leaf counts and the sum of non-leaf counts plus 1 (current node)

Implementation

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

class Solution:
    def solve(self, n):
        if not n:
            return 0, 0
        
        if not n.left and not n.right:
            return 1, 0
        
        left, right = self.solve(n.left), self.solve(n.right)
        return left[0] + right[0], 1 + left[1] + right[1]

# Create the binary tree
ob = Solution()
root = TreeNode(6)
root.left = TreeNode(2)
root.right = TreeNode(6)
root.right.left = TreeNode(10)
root.right.right = TreeNode(2)

print(ob.solve(root))

(3, 2)

How It Works

The algorithm uses post-order traversal to visit each node and count leaves and non-leaves:

1. Base case: If node is null, return (0, 0)
2. Leaf case: If node has no children, return (1, 0) indicating one leaf
3. Recursive case: Count leaves and non-leaves in both subtrees, then combine results

The tuple format (leaf_count, non_leaf_count) makes it easy to track both values simultaneously during the traversal.

Conclusion

This recursive solution efficiently counts leaf and non-leaf nodes in O(n) time by visiting each node once. The algorithm returns a tuple with leaf count first and non-leaf count second.