5 Best Ways to Find the Sum of All Nodes in a Binary Tree Using Python

💡 Problem Formulation: When working with binary trees in Python, a common task is to compute the sum of all the node values. The problem statement is straightforward: given the root of a binary tree, return the sum of the nodes’ values. For instance, a binary tree with nodes 1, 2, and 3 should return a sum of 6.

Method 1: Recursive Depth-First Search (DFS)

Recursive DFS is an elegant approach to traverse a binary tree and sum all node values. It involves visiting each node exactly once and adding up all the values. The function calls itself recursively for left and right subtrees, starting at the root and traversing down. This is a highly readable solution for tree-based problems.

Here’s an example:

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

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

# Example Tree
root = TreeNode(1)
root.left = TreeNode(2)
root.right = TreeNode(3)
print(sum_of_nodes(root))

Output: 6

The code defines a simple recursive function sum_of_nodes() that checks if a node is null (base case) returning 0. For non-null nodes, it returns the value of the node plus the sum of the values from the left and right subtrees. When run with the example tree, it outputs 6, summing the values 1, 2, and 3.

Method 2: Iterative DFS Using Stack

An iterative version of DFS uses a stack data structure to simulate the recursive calls of the first method. It is often used when a non-recursive solution is preferred or when dealing with very deep trees that could cause a stack overflow with recursion.

Here’s an example:

def sum_of_nodes_iterative(root):
    if not root:
        return 0
    stack, node_sum = [root], 0
    while stack:
        node = stack.pop()
        node_sum += node.val
        if node.right:
            stack.append(node.right)
        if node.left:
            stack.append(node.left)
    return node_sum

print(sum_of_nodes_iterative(root))  # Reusing the TreeNode class and root from Method 1

Output: 6

This method uses a stack to keep track of nodes yet to visit while performing DFS on the tree. The sum_of_nodes_iterative() function processes the nodes, summing their values until no more nodes are left on the stack. The final result is the sum of all node values.

Method 3: Iterative Breadth-First Search (BFS) Using Queue

Breadth-First Search (BFS) is another tree traversal method that uses a queue to visit nodes level-by-level. This approach might be more preferable in some scenarios, particularly when minimizing the depth of the traversal matters and for tasks associated with the level of nodes.

Here’s an example:

from collections import deque

def sum_of_nodes_bfs(root):
    if not root:
        return 0
    queue, node_sum = deque([root]), 0
    while queue:
        node = queue.popleft()
        node_sum += node.val
        if node.left:
            queue.append(node.left)
        if node.right:
            queue.append(node.right)
    return node_sum

print(sum_of_nodes_bfs(root))  # Reusing the TreeNode class and root from Method 1

Output: 6

The sum_of_nodes_bfs() function systematically processes nodes by their level, using a queue to implement BFS traversal. Similar to DFS, it aggregates the values until all nodes have been visited, resulting in the sum of values.

Method 4: Divide and Conquer

The divide and conquer strategy is a paradigm that divides a problem into subproblems, conquers them recursively, and combines their results. This way, we divide our tree into subtrees and solve for each before combining results.

Here’s an example:

def sum_of_nodes_divide_conquer(root):
    return sum_of_nodes(root)  # Utilizing the recursive DFS function from Method 1

print(sum_of_nodes_divide_conquer(root))

Output: 6

This example reuses the sum_of_nodes() function from Method 1 as it inherently follows the divide and conquer approach by recursively solving for the sum of left and right subtrees. The solutions are then combined by summing the node’s value with the received sums from subtree calls.

Bonus One-Liner Method 5: Pythonic Recursive DFS

For Python enthusiasts who admire Python’s conciseness, a one-liner variant of the recursive DFS method is to use the walrus operator (:=) introduced in Python 3.8. This method captures and returns the sum in a single expression.

Here’s an example:

sum_of_nodes = lambda root: root.val + sum_of_nodes(root.left) + sum_of_nodes(root.right) if root else 0

print(sum_of_nodes(root))

Output: 6

This compact lambda function is a one-liner equivalent of the recursive DFS approach in Method 1. It evaluates the existence of the root, then computes the sum of its value with the recursive calls on the left and right subtrees.

Summary/Discussion

• Method 1: Recursive DFS. Intuitive. Could lead to stack overflow for very deep trees.
• Method 2: Iterative DFS using stack. Non-recursive. Uses more memory for manually managing stack.
• Method 3: Iterative BFS using queue. Level-based traversal. Uses more memory for queue management, slightly less efficient for sum calculation.
• Method 4: Divide and Conquer. Conceptual approach. Internally similar to recursive DFS.
• Method 5: Pythonic Recursive DFS One-Liner. Elegant and concise. Requires Python 3.8+, less readable for beginners.