Replace node with sum of all greater node in binary search tree

Today, we will see a problem, solution to which uses traversal of tree in a bit different manner. Problem statement is : Replace a node with sum of all nodes which are greater than it. First question which comes to mind is : what should be done with rightmost child of tree? Should it remain as such or replaced with zero as there is node greater than it? Based on clarification, we can design our solution. In our case, for simplicity, we will replace right most node with zero.

Example : Original tree

replace with sum of greater nodes

After replace nodes with sum of all greater nodes.

replace node with sum of greater nodes

How can we solve it? Easiest way is to do inorder traversal of BST and store that into an array. Now, starting from end of array, replace element of array with sum of all elements on the right. With bit of trick, this algorithm works in O(N) time, however, uses two traversals of tree and O(N) extra space.

We know that inorder traversal of BST gives nodes in sorted in ascending order. Since, we need to visit greater nodes before the node in order to calculate their sum, can we just have nodes of BST visited in sorted but in descending order.

How about reversing the inorder traversal itself? In this, traverse right subtree first, then node and at last left subtree. When right subtree is traversed for a node, we are actually traversing all node which are greater that node, all we need to do is return sum from that traversal. While traversing left subtree, we have to just propagate sum of right subtree and node.

Solution can be implemented using recursion with base case being right most node of tree. If we reached right most tree,  replace it with zero and return the value of the node to parent.

Implementation of replacing a node with sum of all greater nodes

Complexity of algorithm to replace node with sum of all nodes greater than it, is O(N) with no extra space.