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
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
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
Java
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
}
}
Python
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, wherehis the height of the tree (O(log n)for balanced trees,O(n)for skewed trees).