从二叉搜索树到更大和树 : BST 中序遍历的经典运用

题目描述

这是 LeetCode 上的 「1038. 从二叉搜索树到更大和树」 ，难度为 「中等」。

Tag : 「BST」、「中序遍历」

给定一个二叉搜索树 root (BST)，请将它的每个节点的值替换成树中大于或者等于该节点值的所有节点值之和。

提醒一下，二叉搜索树满足下列约束条件：

节点的左子树仅包含键小于节点键的节点。
节点的右子树仅包含键大于节点键的节点。
左右子树也必须是二叉搜索树。

示例 1： alt

输入：[4,1,6,0,2,5,7,null,null,null,3,null,null,null,8]

输出：[30,36,21,36,35,26,15,null,null,null,33,null,null,null,8]

示例 2：

输入：root = [0,null,1]

输出：[1,null,1]

提示：

树中的节点数在范围内。
树中的所有值均不重复。

中序遍历

利用 「BST 的中序遍历是有序」 的特性，我们可以通过两次遍历 BST 来求解问题。

首先，通过一次遍历，计算出整棵树的节点总和 tot，然后在中序遍历过程中，不断对 tot 进行更新，将其作为当前未遍历到的节点的总和，用于给当前节点赋值。

假设当前遍历到的节点为 x（起始节点值为 t），那么将节点更新为当前节点 tot 后，更新 tot = tot - t。

这是常规的中序遍历做法，更进一步，如果将其中序遍历的顺序进行翻转（从「左中右」调整为「右中左」），则可实现一次遍历。

Java 代码：

class Solution {
    int tot = 0;
    public TreeNode bstToGst(TreeNode root) {
        dfs(root);
        return root;
    }
    void dfs(TreeNode root) {
        if (root == null) return ;
        dfs(root.right);
        tot += root.val;
        root.val = tot;
        dfs(root.left);
    }
}

C++ 代码：

class Solution {
public:
    int tot = 0;
    TreeNode* bstToGst(TreeNode* root) {
        dfs(root);
        return root;
    }
    void dfs(TreeNode* root) {
        if (root == nullptr) return;
        dfs(root->right);
        tot += root->val;
        root->val = tot;
        dfs(root->left);
    }
};

Python 代码：

class Solution:
    def bstToGst(self, root: TreeNode) -> TreeNode:
        tot = 0
        def dfs(root):
            nonlocal tot
            if not root: return
            dfs(root.right)
            tot += root.val
            root.val = tot
            dfs(root.left)
        dfs(root)
        return root

TypeScript 代码：

function bstToGst(root: TreeNode | null): TreeNode | null {
    let tot = 0;
    const dfs = function(root: TreeNode | null): void {
        if (!root) return ;
        dfs(root.right);
        tot += root.val;
        root.val = tot;
        dfs(root.left);
    }
    dfs(root);
    return root;
};