Program to find sum of all elements of a tree in Python

A binary tree is a hierarchical data structure where each node has at most two children. To find the sum of all elements in a binary tree, we can use recursive traversal to visit each node and accumulate their values.

Given a binary tree like this:

The sum would be: 2 + 4 + 3 + 5 = 14

Approach

We'll use a recursive approach to traverse the tree:

• Start with the current node's value

• Recursively add the sum of the left subtree (if it exists)

• Recursively add the sum of the right subtree (if it exists)

• Return the total sum

Implementation

First, let's define the tree node structure and implement the solution:

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

class Solution:
    def recurse(self, node):
        val = node.val
        if node.left:
            val += self.recurse(node.left)
        if node.right:
            val += self.recurse(node.right)
        return val
    
    def solve(self, root):
        if not root:
            return 0
        return self.recurse(root)

# Create the binary tree
root = TreeNode(2)
root.right = TreeNode(4)
root.right.left = TreeNode(3)
root.right.right = TreeNode(5)

# Find sum of all elements
ob = Solution()
result = ob.solve(root)
print(f"Sum of all elements: {result}")

Sum of all elements: 14

Alternative Approach - Simplified Version

We can simplify this into a single function:

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

def sum_tree(root):
    if not root:
        return 0
    
    return root.val + sum_tree(root.left) + sum_tree(root.right)

# Create the same binary tree
root = TreeNode(2)
root.right = TreeNode(4)
root.right.left = TreeNode(3)
root.right.right = TreeNode(5)

# Calculate sum
total_sum = sum_tree(root)
print(f"Sum using simplified approach: {total_sum}")

Sum using simplified approach: 14

How It Works

The recursive algorithm works as follows:

• Base case: If the current node is None (empty), return 0

• Recursive case: Return current node's value + sum of left subtree + sum of right subtree

• The recursion naturally handles all nodes in the tree

Time and Space Complexity

• Time Complexity: O(n), where n is the number of nodes (we visit each node once)

• Space Complexity: O(h), where h is the height of the tree (due to recursion stack)

Conclusion

Finding the sum of all elements in a binary tree is efficiently solved using recursion. The algorithm visits each node exactly once and accumulates their values, making it optimal with O(n) time complexity.