Python Program to Find the Sum of all Nodes in a Tree

Finding the sum of all nodes in a tree is a fundamental operation in tree data structures. This can be accomplished using a Tree_structure class with methods to add nodes and calculate the total sum recursively.

Tree Structure Implementation

The following example demonstrates a complete tree implementation with sum calculation ?

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

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

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

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

    def summation_nodes(self):
        sum_val = self.key
        for child in self.children:
            sum_val = sum_val + child.summation_nodes()
        return sum_val

# Create a simple tree structure
tree = Tree_structure(10)
tree.add_values(Tree_structure(5))
tree.add_values(Tree_structure(15))
tree.children[0].add_values(Tree_structure(3))
tree.children[0].add_values(Tree_structure(7))

print("Sum of all nodes:", tree.summation_nodes())
Sum of all nodes: 40

Interactive Tree Builder

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

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

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

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

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

    def summation_nodes(self):
        sum_val = self.key
        for child in self.children:
            sum_val = sum_val + child.summation_nodes()
        return sum_val

tree = None

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

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

    operation = my_input[0].strip().lower()
    if operation == 'add':
        data = int(my_input[1])
        newNode = Tree_structure(data)
        sub_op = my_input[2].strip().lower()
        if sub_op == 'at':
            tree = newNode
        elif sub_op == 'below':
            my_pos = my_input[3].strip().lower()
            key = int(my_pos)
            ref_node = None
            if tree is not None:
                ref_node = tree.search_val(key)
            if ref_node is None:
                print('No such key exists')
                continue
            ref_node.add_values(newNode)

    elif operation == 'summation':
        if tree is None:
            print('The tree is empty')
        else:
            summation_val = tree.summation_nodes()
            print('Sum of all the nodes is : {}'.format(summation_val))

    elif operation == 'quit':
        break

How the Sum Calculation Works

The summation_nodes() method uses recursion to calculate the total sum:

  1. Start with the current node's value
  2. Recursively add the sum of each child subtree
  3. Return the total sum

Key Features

  • Recursive sum calculation: Each node calculates its subtree sum
  • Dynamic tree building: Add nodes interactively at runtime
  • Node search: Find specific nodes to add children below them
  • Error handling: Checks for empty trees and invalid keys

Sample Output

Menu (no duplicate keys allowed)
add <data> at root
add <data> below <data>
summation
quit
What would you like to do? add 56 at root
What would you like to do? add 45 below 56
What would you like to do? add 23 below 56
What would you like to do? summation
Sum of all the nodes is : 124
What would you like to do? quit

Conclusion

The recursive approach efficiently calculates the sum of all nodes by traversing each subtree once. This tree implementation provides a flexible foundation for various tree operations and demonstrates fundamental recursion concepts.

Updated on: 2026-03-25T18:45:15+05:30

479 Views

Kickstart Your Career

Get certified by completing the course

Get Started
Advertisements