How to Convert BST to Greater Tree?

Given a Binary Search Tree (BST), convert it to a Greater Tree such that every key of the original BST is changed to the original key plus sum of all keys greater than the original key in BST.

Example:

Input: The root of a Binary Search Tree like this:

              5
            /   \
           2     13

Output: The root of a Greater Tree like this:

             18
            /   \
          20     13

Post-order Traversal Recursion

The post-order traversal of a BST (Binary Search Tree) gives the reverse sorting order. Therefore, equivalently, we scan the array from biggest to smallest and we have a counter to sum up the nodes we have visited – then update the nodes along the way.

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

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    TreeNode* convertBST(TreeNode* root) {
        if (root == NULL) return root;
        convertBST(root->right);
        sum += root->val;
        root->val = sum;
        convertBST(root->left);
        return root;
    }
private:
    int sum = 0;
};

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    TreeNode* convertBST(TreeNode* root) {
        if (root == NULL) return root;
        convertBST(root->right);
        sum += root->val;
        root->val = sum;
        convertBST(root->left);
        return root;
    }
private:
    int sum = 0;
};

The recursion implementation illustrates the idea with minimal amount of code – the compiler generates and maintains the stack automatically at runtime.

Post-order Traversal Iterative Approach with Stack

With a manual stack, we can implement the above idea with iterative approach. First push all the right nodes of the BST to the stack, pop one by one, increment the counter and push the left nodes to the stack until the stack is empty and the node is NULL.

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

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    TreeNode* convertBST(TreeNode* root) {
        int sum = 0;
        TreeNode* node = root;
        stack<TreeNode*> st;
        while ((st.size() > 0) || node != NULL) {
            while (node != NULL) {
                st.push(node);
                node = node->right;    
            }
            node = st.top();
            st.pop();
            sum += node->val;
            node->val = sum;
            node = node->left;
        }
        return root;
    }
};

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    TreeNode* convertBST(TreeNode* root) {
        int sum = 0;
        TreeNode* node = root;
        stack<TreeNode*> st;
        while ((st.size() > 0) || node != NULL) {
            while (node != NULL) {
                st.push(node);
                node = node->right;    
            }
            node = st.top();
            st.pop();
            sum += node->val;
            node->val = sum;
            node = node->left;
        }
        return root;
    }
};

Both C++ approaches are O(N) time and space complexity – as we need to visit all the N nodes and the stack size is N depth the worse case – when the BST is degenerated into a linked list.

–EOF (The Ultimate Computing & Technology Blog) —