Replace each node in BST with subtree sum and its value

Problem

Given a binary search tree (BST), the task is to modify it such that each node’s value is replaced with the sum of the values of:

all nodes in its left subtree,
all nodes in its right subtree,
its own value.

Examples

Example 1

graph LR
    subgraph Input
        A1(6)
        A1 --- B1(3)
        A1 --- C1(8)
        B1 --- D1(1)
        B1 --- E1(4)
        C1 ~~~ N1:::hidden
        C1 --- F1(10)
    end

    subgraph Output
        A2(32)
        A2 --- B2(8)
        A2 --- C2(18)
        B2 --- D2(1)
        B2 --- E2(4)
        C2 ~~~ N2:::hidden
        C2 --- F2(10)
    end

Input --> Output
classDef hidden display:none

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19


Input:  
	 6
	/ \
   3   8
  / \    \
 1   4    10

Output:  
	 32
	/  \
   8    18
  / \     \
 1   4     10

Explanation:  
- Node 6 is replaced with 6 + (3+1+4) + (8+10) = 32  
- Node 3 is replaced with 3 + (1) + (4) = 8  
- Node 8 is replaced with 8 + (0) + (10) = 18  
- Node 1, 4, and 10 remain unchanged as they have no children.

Example 2

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18


Input:  
	 5
	/ \
   2   7
	  /
	 6

Output:  
	 20
	/  \
   2   13
	  /
	 6

Explanation:  
- Node 5 is replaced with 5 + 2 + (7+6) = 20  
- Node 7 is replaced with 7 + (6) + 0 = 13  
- Node 2 and 6 remain unchanged as they have no children.

Solution

Method 1 - Post order Traversal

To update the root node of the given tree, we need the sum of all nodes in the left subtree and the right subtree. This makes it clear that the left and right subtrees must be processed first, following a typical post-order traversal approach. Solving the problem at the root depends on solving it at the left and right subtree levels first, which naturally leads to a recursive implementation.

At the leaf nodes, there are no children, so their sum remains unchanged. It is often assumed to keep leaf nodes as they are, but it’s advisable to confirm this assumption with the interviewer before proceeding with the solution.

In a post-order traversal, the process involves first calculating the sum of the left and right subtrees, then updating the root node’s value to the total sum. To grasp the flow and understand what needs to be done, working through an example manually is highly recommended before coding. For instance, we can take a sample tree and replace each node with the sum of its children to solidify the concept.

Approach

Dry Run

graph TD;
	F(6)
	F --- C(3)
	F --- H(8)
	C --- A(1)
	C --- D(4)
	H ~~~ N1:::hidden
	H --- J(10)

classDef hidden display:none
classDef blue fill:#1E90FF,stroke:#000,stroke-width:1px,color:#fff;

We begin at the root node (6). To update its value, we first compute the sum of values in its left and right subtrees. Hence, we proceed down to its left child node (3).

graph TD;
	F(6):::blue
	F --- C(3)
	F --- H(8)
	C --- A(1)
	C --- D(4)
	H ~~~ N1:::hidden
	H --- J(10)

classDef hidden display:none
classDef blue fill:#1E90FF,stroke:#000,stroke-width:1px,color:#fff;

At node (3), the same process applies. We need the subtree sums, so we move further down to its left child, node (1).

graph TD;
	F(6)
	F --- C(3):::blue
	F --- H(8)
	C --- A(1)
	C --- D(4)
	H ~~~ N1:::hidden
	H --- J(10)

classDef hidden display:none
classDef blue fill:#1E90FF,stroke:#000,stroke-width:1px,color:#fff;

Since node (1) is a leaf node with no children, its value remains unchanged. We return its value (1) back to parent node (3).

graph TD;
	F(6)
	F --- C(3)
	F --- H(8)
	C --- A(1):::blue
	C --- D(4)
	H ~~~ N1:::hidden
	H --- J(10)

classDef hidden display:none
classDef blue fill:#1E90FF,stroke:#000,stroke-width:1px,color:#fff;

Back at node (3), we now process its right subtree. We move to node (4).

graph TD;
	F(6)
	F --- C(3):::blue
	F --- H(8)
	C --- A(1)
	C --- D(4)
	H ~~~ N1:::hidden
	H --- J(10)

classDef hidden display:none
classDef blue fill:#1E90FF,stroke:#000,stroke-width:1px,color:#fff;

Like node (1), node (4) is a leaf, so its value remains unchanged. We return its value (4) back to parent node (3).

graph TD;
	F(6)
	F --- C(3)
	F --- H(8)
	C --- A(1)
	C --- D(4):::blue
	H ~~~ N1:::hidden
	H --- J(10)

classDef hidden display:none
classDef blue fill:#1E90FF,stroke:#000,stroke-width:1px,color:#fff;

At this point, both left and right subtree sums of node (3) are available. We calculate the new value for node (3) as 1 + 3 + 4 = 8 and replace its value with this sum.

graph TD;
	F(6)
	F --- C(8):::blue
	F --- H(8)
	C --- A(1)
	C --- D(4)
	H ~~~ N1:::hidden
	H --- J(10)

classDef hidden display:none
classDef blue fill:#1E90FF,stroke:#000,stroke-width:1px,color:#fff;

With the left subtree of the root node (6) fully processed, we move to parent node(6).

graph TD;
	F(6):::blue
	F --- C(8)
	F --- H(8)
	C --- A(1)
	C --- D(4)
	H ~~~ N1:::hidden
	H --- J(10)

classDef hidden display:none
classDef blue fill:#1E90FF,stroke:#000,stroke-width:1px,color:#fff;

And, proceed to its right subtree, moving to node (8).

graph TD;
	F(6)
	F --- C(8)
	F --- H(8):::blue
	C --- A(1)
	C --- D(4)
	H ~~~ N1:::hidden
	H --- J(10)

classDef hidden display:none
classDef blue fill:#1E90FF,stroke:#000,stroke-width:1px,color:#fff;

Node (8) has no left subtree, so we directly process its right subtree, moving down to node (10).

graph TD;
	F(6)
	F --- C(8)
	F --- H(8)
	C --- A(1)
	C --- D(4)
	H ~~~ N1:::hidden
	H --- J(10):::blue

classDef hidden display:none
classDef blue fill:#1E90FF,stroke:#000,stroke-width:1px,color:#fff;

Since node (10) is a leaf, its value remains unchanged, and we return its value back to parent node (8). Now, both left and right subtree sums of node (8) are available. We calculate the new value for node (8) as 8 + 0 + 10 = 18 and replace its value with this sum.

graph TD;
	F(6)
	F --- C(8)
	F --- H(18):::blue
	C --- A(1)
	C --- D(4)
	H ~~~ N1:::hidden
	H --- J(10)

classDef hidden display:none
classDef blue fill:#1E90FF,stroke:#000,stroke-width:1px,color:#fff;

Finally, at the root node (6), with sums of both its left and right subtrees computed, we update its value as 6 + 8 + 18 = 32.

graph TD;
	F(32):::blue
	F --- C(8)
	F --- H(18)
	C --- A(1)
	C --- D(4)
	H ~~~ N1:::hidden
	H --- J(10)

classDef hidden display:none
classDef blue fill:#1E90FF,stroke:#000,stroke-width:1px,color:#fff;

This process ensures every node in the tree is updated to the sum of its subtree values and its original value.

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47


class Solution {

    // Function to modify the tree
    public int modifyTree(TreeNode node) {
        if (node == null) {
            return 0;
        }

        // Post-order traversal: Left, Right, Node
        int leftSum = modifyTree(node.left);
        int rightSum = modifyTree(node.right);

        // Calculate new value for the current node
        int originalValue = node.val;
        node.val = originalValue + leftSum + rightSum;

        // Return the cumulative sum to parent node
        return node.val;
    }

    // Utility function to print the tree (In-order traversal)
    public void printTree(TreeNode node) {
        if (node == null) {
            return;
        }
        printTree(node.left);
        System.out.print(node.val + " ");
        printTree(node.right);
    }

    public static void main(String[] args) {
        Solution obj = new Solution();
        // Create the original BST
        TreeNode root = new TreeNode(6);
        root.left = new TreeNode(3);
        root.right = new TreeNode(8);
        root.left.left = new TreeNode(1);
        root.left.right = new TreeNode(4);
        root.right.right = new TreeNode(10);

        // Modify tree
        obj.modifyTree(root);

        // Print modified tree
        obj.printTree(root); // Output: 1 8 4 32 18 10
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39


class Solution:
    def modifyTree(self, node):
        if not node:
            return 0
        
        # Post-order traversal: Left, Right, Node
        left_sum = self.modifyTree(node.left)
        right_sum = self.modifyTree(node.right)
        
        # Calculate new value for the current node
        original_value = node.val
        node.val = original_value + left_sum + right_sum
        
        # Return the cumulative sum to parent node
        return node.val

    def printTree(self, node):
        if not node:
            return
        self.printTree(node.left)
        print(node.val, end=" ")
        self.printTree(node.right)

# Example usage
if __name__ == "__main__":
    # Create the original BST
    root = TreeNode(6)
    root.left = TreeNode(3)
    root.right = TreeNode(8)
    root.left.left = TreeNode(1)
    root.left.right = TreeNode(4)
    root.right.right = TreeNode(10)
    
    # Modify tree
    solution = Solution()
    solution.modifyTree(root)
    
    # Print modified tree
    solution.printTree(root)  # Output: 1 8 4 32 18 10

Complexity

⏰ Time complexity: O(n). Each node is visited once, and subtree sums are calculated in constant time.
🧺 Space complexity: O(h). Requires space for the recursive stack, where h is the height of the tree (O(log n) for balanced trees, O(n) for skewed trees).

Problem#

Examples#

Example 1#

Example 2#

Solution#

Method 1 - Post order Traversal#

Approach#

Dry Run#

Code#

Complexity#