P141 将二叉搜索树转为较大树

描述#

这个题目说的是，给你一棵二叉搜索树，你要把它转换为一棵较大树(Greater Tree)。转换方法是，将二叉搜索树中每个节点值更新为它与所有大于它的节点值之和。

比如说，给你的二叉搜索树是：
  4
/   \

2     5
/ \     

1   3     6
转换后的较大树为：
  15
/    \

20    11
/  \     

21  18     6
可以看出，较大树的结构和原始二叉搜索树一样，而节点值等于原始二叉树中该节点的值加上所有比它大的节点值。

代码#

public class AlgoCasts {

  public class TreeNode {
    int val;
    TreeNode left;
    TreeNode right;

    TreeNode(int x) {
      val = x;
    }
  }

  private int dfs(TreeNode root, int sum) {
    if (root == null) return sum;
    int cur = dfs(root.right, sum);
    root.val += cur;
    return dfs(root.left, root.val);
  }

  // Time: O(n), Space: O(n)
  public TreeNode convertBSTRecursive(TreeNode root) {
    dfs(root, 0);
    return root;
  }

  // Time: O(n), Space: O(n)
  public TreeNode convertBSTIterative(TreeNode root) {
    Stack<TreeNode> st = new Stack<>();
    TreeNode cur = root;
    int sum = 0;

    while (cur != null || !st.empty()) {
      while (cur != null) {
        st.push(cur);
        cur = cur.right;
      }
      cur = st.pop();
      cur.val += sum;
      sum = cur.val;
      cur = cur.left;
    }
    return root;
  }

}