LeetCode 538. 把二叉搜索树转换为累加树

Offer 驾到，掘友接招！我正在参与2022春招打卡活动，点击查看活动详情。

题目：给定二叉搜索树的根节点，树中所有节点的值都不同，要求将其转为累加树，使得每个节点的新值大于等于原树中大于等于该节点之和。

解题思路

做这题真被自己傻到了，且看解题思路。

以下面这颗树为例：

根节点4的值应该是其自身节点和右子树所有节点之和，而左子树1的新值应该是除了1左子树的其余节点之和，至此得到一个规律：

每个节点的新值为：父节点的新值 + 自身节点 + 右子树所有节点。

根据此规律，我们遍历整颗二叉树，每计算一个节点都计算以这个节点为根节点的树的全部节点之和，之和用此值减去左子树的值再加上父节点的值即为最终答案，可得代码如下：

public TreeNode convertBST(TreeNode root) {
    return get(root, 0);
}

public TreeNode get(TreeNode root, int sum){
    if(root==null) return null;
    root.val =  sum + getTreeSum(root) - getTreeSum(root.left);
    root.left = get(root.left, root.val);
    root.right = get(root.right, sum);
    return root;
}

public int getTreeSum(TreeNode root){
    int treeSum = 0;
    if (root == null) return treeSum;

    Stack<TreeNode> stack = new Stack<>();
    stack.push(root);
    while (!stack.isEmpty()){
        TreeNode pop = stack.pop();
        treeSum += pop.val;
        if(pop.left!=null) stack.push(pop.left);
        if(pop.right!=null) stack.push(pop.right);
    }
    return treeSum;
}

最终耗时33ms，超过5%的小伙伴。想想代码好像没问题，于是想着优化代码：

此处注意，在计算根节点的整颗子树之和的时候，我们还计算了以下面所有节点为根节点时的子树之和，那么在再次计算子节点时候，出现了大量的重复计算，这显然是很耗时的，实际上我们可以在计算根节点的时候就将所有子树节点的值存起来，之后再次访问子节点的时候只需要先查询一下是否有即可，这样计算树的所有节点之和只计算了一次，可得代码如下：

public TreeNode convertBST(TreeNode root) {
    HashMap<TreeNode, Integer> map = new HashMap<>();
    return get(root, 0, map);
}

public TreeNode get(TreeNode root, int sum, HashMap<TreeNode, Integer> map){
    if(root==null) return null;
    root.val =  sum + getTreeSum2(root, map) - getTreeSum2(root.left, map);
    root.left = get(root.left, root.val, map);
    root.right = get(root.right, sum, map);
    return root;
}

public int getTreeSum2(TreeNode root, HashMap<TreeNode, Integer> map){
    if(root == null) return 0;
    if(map.containsKey(root)) return map.get(root);
    int sum = root.val;
    sum += getTreeSum2(root.left, map);
    sum += getTreeSum2(root.right, map);
    map.put(root, sum);
    return sum;
}

最终耗时3ms，想想优化了十倍还沾沾自喜。。。最后一看还是只超过了5%的小伙伴。。。

我们再看上面那棵树，有没有发现从最右边节点开始，每个节点的新值都是以下规律：

1. 右子树节点 = 自身 + 右子树和
2. 右子树节点父节点 = 自身 + 右子树和
3. 左子树节点 = 自身 + 父节点的右子树节点 + 父节点

这显然就是一个反中序排列，我们只需要对二叉树进行一次反中序排列之后累加和即可~

private int sum = 0;

public TreeNode convertBST(TreeNode root) {
    if(root!=null){
        convertBST(root.right);
        sum += root.val;
        root.val = sum;
        convertBST(root.left);
    }
    return root;
}

最终耗时0ms。