Program to find sum each of the diagonal path elements in a binary tree in Python

In this problem, we need to find the sum of elements along each diagonal path in a binary tree. A diagonal path moves from top-left to bottom-right, where moving to the right child continues the same diagonal, while moving to the left child starts a new diagonal level.

Looking at the example tree structure:

Algorithm Approach

The key insight is that each diagonal can be identified by how many "left moves" we made to reach it. Moving right stays on the same diagonal, while moving left increases the diagonal level by 1.

Implementation

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

class Solution:
    def solve(self, root):
        output = []
        
        def traverse(node, numLeft, output):
            if not node:
                return
            
            # If this is a new diagonal level, append the node's value
            if numLeft >= len(output):
                output.append(node.data)
            else:
                # Add to existing diagonal sum
                output[numLeft] += node.data
            
            # Left child increases diagonal level
            if node.left:
                traverse(node.left, numLeft + 1, output)
            
            # Right child stays on same diagonal level
            if node.right:
                traverse(node.right, numLeft, output)
        
        traverse(root, 0, output)
        return output

# Create the example tree
ob = Solution()
root = TreeNode(12)
root.left = TreeNode(8)
root.right = TreeNode(15)
root.left.left = TreeNode(3)
root.left.right = TreeNode(10)

print(ob.solve(root))

[27, 18, 3]

How It Works

The algorithm performs a depth-first traversal while tracking the diagonal level:

• numLeft parameter tracks how many left moves we've made

• When visiting a node, we either start a new diagonal or add to an existing one

• Moving to right child keeps the same diagonal level

• Moving to left child increases diagonal level by 1

Trace Through Example

1. Start at root (12), diagonal 0: output = [12]

2. Go right to 15, same diagonal 0: output = [27]

3. Go left to 8, diagonal 1: output = [27, 8]

4. Go right to 10, same diagonal 1: output = [27, 18]

5. Go left to 3, diagonal 2: output = [27, 18, 3]

Conclusion

This solution efficiently calculates diagonal sums by tracking the "left depth" during tree traversal. The time complexity is O(n) where n is the number of nodes, and space complexity is O(h) for the recursion stack where h is the tree height.