538 Convert BST to Greater Tree 把二叉搜索樹轉換為累加樹

Description:
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

題目描述:
給定一個二叉搜索樹（Binary Search Tree），把它轉換成為累加樹（Greater Tree)，使得每個節點的值是原來的節點值加上所有大于它的節點值之和。

示例 :

例如：

輸入: 二叉搜索樹:

              5
            /   \
           2     13

輸出: 轉換為累加樹:

             18
            /   \
          20     13

思路:

1. 先計算所有數字之和, 然后按照先序遍歷給各個結點賦值
2. 按照右 -> 根 -> 左的順序記錄結點的值并累加
   可以用遞歸和迭代解決
   時間復雜度O(n), 空間復雜度O(n)

代碼:
C++:

/**
 * 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 pre = 0;
        order(root, pre);
        return root;
    }
private:
    void order(TreeNode* root, int &pre) 
    {
        if (!root) return;
        stack<TreeNode*> s;
        while (root or s.size()) 
        {
            if (root) 
            {
                s.push(root);
                root = root -> right;
            } 
            else 
            {
                root = s.top();
                s.pop();
                root -> val += pre;
                pre = root -> val;
                root = root -> left;
            }
        }
    }
};

Java:

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
class Solution {
    private int pre = 0;
    public TreeNode convertBST(TreeNode root) {
        order(root);
        return root;
    }

    private void order(TreeNode root) {
        if (root == null) return;
        order(root.right);
        root.val += pre;
        pre = root.val;
        order(root.left);
    }

}

Python:

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution:
    def convertBST(self, root: TreeNode) -> TreeNode:
        s, l, p, pre = 0, [], root, 0
        def dfs(root: TreeNode) -> int:
            if not root:
                return 0
            return root.val + dfs(root.left) + dfs(root.right)
        s = dfs(root)
        while p or l:
            if p:
                l.append(p)
                p = p.left
            else:
                p = l.pop()
                pre = p.val
                p.val = s
                s -= pre
                p = p.right
        return root