Python Program to Find the Sum of All Nodes in a Binary Tree

When working with binary trees, finding the sum of all nodes is a common operation. We can implement this using a tree class that contains methods to build the tree and calculate the sum recursively. Each node stores data and maintains references to its children.

Below is a demonstration of the same −

Tree Class Implementation

First, let's create a tree structure class with methods for adding nodes and calculating the sum ?

class Tree_struct:
    def __init__(self, data=None):
        self.key = data
        self.children = []

    def set_root(self, data):
        self.key = data

    def add_node(self, node):
        self.children.append(node)

    def search_node(self, key):
        if self.key == key:
            return self
        for child in self.children:
            temp = child.search_node(key)
            if temp is not None:
                return temp
        return None

    def sum_node(self):
        my_summation = self.key
        for child in self.children:
            my_summation = my_summation + child.sum_node()
        return my_summation

# Create a simple tree for demonstration
root = Tree_struct(10)
child1 = Tree_struct(5)
child2 = Tree_struct(15)
child3 = Tree_struct(3)

root.add_node(child1)
root.add_node(child2)
child1.add_node(child3)

print("Tree structure created:")
print("Root: 10")
print("Children of 10: 5, 15") 
print("Child of 5: 3")
print(f"Sum of all nodes: {root.sum_node()}")

Tree structure created:
Root: 10
Children of 10: 5, 15
Child of 5: 3
Sum of all nodes: 33

Interactive Tree Builder

Here's a complete program that allows users to build a tree interactively and calculate the sum ?

class Tree_struct:
    def __init__(self, data=None):
        self.key = data
        self.children = []

    def set_root(self, data):
        self.key = data

    def add_node(self, node):
        self.children.append(node)

    def search_node(self, key):
        if self.key == key:
            return self
        for child in self.children:
            temp = child.search_node(key)
            if temp is not None:
                return temp
        return None

    def sum_node(self):
        my_summation = self.key
        for child in self.children:
            my_summation = my_summation + child.sum_node()
        return my_summation

my_instance = None

print('Menu (assume no duplicate keys)')
print('add <data> at root')
print('add <data> below <data>')
print('sum')
print('quit')

while True:
    my_input = input('What operation would you do ? ').split()

    operation = my_input[0].strip().lower()
    if operation == 'add':
        data = int(my_input[1])
        new_node = Tree_struct(data)
        suboperation = my_input[2].strip().lower()
        if suboperation == 'at':
            my_instance = new_node
        elif suboperation == 'below':
            position = my_input[3].strip().lower()
            key = int(position)
            ref_node = None
            if my_instance is not None:
                ref_node = my_instance.search_node(key)
            if ref_node is None:
                print('No such key')
                continue
            ref_node.add_node(new_node)

    elif operation == 'sum':
        if my_instance is None:
            print('The tree is empty')
        else:
            my_summation = my_instance.sum_node()
            print('Sum of all nodes is: {}'.format(my_summation))

    elif operation == 'quit':
        break

Sample Output

Menu (assume no duplicate keys)
add <data> at root
add <data> below <data>
sum
quit
What operation would you do ? add 5 at root
What operation would you do ? add 7 below 5
What operation would you do ? add 0 below 7
What operation would you do ? sum
Sum of all nodes is: 12
What operation would you do ? quit

How the Sum Calculation Works

The sum_node() method uses recursion to calculate the total sum ?

• Start with the current node's value

• Recursively call sum_node() on each child

• Add all child sums to the current node's value

• Return the total sum

Key Features

• The Tree_struct class represents a general tree (not strictly binary)

• Each node can have multiple children stored in a list

• The search_node() method finds nodes by their key value

• The sum_node() method recursively sums all descendant nodes

• The interactive interface allows dynamic tree construction

Conclusion

This implementation demonstrates how to build a tree data structure and calculate the sum of all nodes using recursion. The sum_node() method efficiently traverses the entire tree and accumulates the sum of all node values.