二叉搜索树题目：修剪二叉搜索树

题目

标题和出处

标题：修剪二叉搜索树

出处：669. 修剪二叉搜索树

难度

3 级

题目描述

要求

给定二叉搜索树的根结点 $\text{root}$ 以及下界 $\text{low}$ 和上界 $\text{high}$，修剪二叉搜索树，使得所有结点的值在 $\text{[low, high]}$ 中。修剪树不应该改变保留在树中的元素的相对结构 （即任何结点的后代仍然应该作为后代）。可以证明，存在唯一的答案。

返回修剪后的二叉搜索树的根结点。注意，根结点可能会根据给定的边界发生改变。

示例

示例 1：

输入：$\text{root = [1,0,2], low = 1, high = 2}$
输出：$\text{[1,null,2]}$

示例 2：

输入：$\text{root = [3,0,4,null,2,null,null,1], low = 1, high = 3}$
输出：$\text{[3,2,null,1]}$

数据范围

• 树中结点数目在范围 ${\text{[1, 10}}^{\text{4}}\text{]}$ 内
• $\text{0}\le \text{Node.val}\le {\text{10}}^{\text{4}}$
• 树中的每个结点值各不相同
• $\text{root}$ 保证是有效的二叉搜索树
• $\text{0}\le \text{low}\le \text{high}\le {\text{10}}^{\text{4}}$

解法一

思路和算法

如果二叉搜索树为空，则不需要修剪，返回空树。

如果二叉搜索树不为空，则首先判断根结点值是否在给定的边界范围内，决定根结点是否保留。

• 如果根结点值大于上界，则根结点的右子树中的所有结点值都大于根结点值，因此都大于上界，根结点和右子树中的所有结点都不保留，修剪后剩余的结点都在根结点的左子树中。

• 如果根结点值小于下界，则根结点的左子树中的所有结点值都小于根结点值，因此都小于下界，根结点和左子树中的所有结点都不保留，修剪后剩余的结点都在根结点的右子树中。

• 如果根结点值在给定的边界范围内，则根结点保留，对左子树和右子树分别修剪。

上述过程是一个递归的过程，递归的终止条件是当前结点为空，此时返回空树。对于其余情况，首先判断根结点是否保留，然后对根结点的左子树和右子树调用递归，并用递归调用的结果更新根结点的左子树和右子树。

代码

class Solution {
    public TreeNode trimBST(TreeNode root, int low, int high) {
        if (root == null) {
            return root;
        }
        if (root.val > high) {
            return trimBST(root.left, low, high);
        }
        if (root.val < low) {
            return trimBST(root.right, low, high);
        }
        root.left = trimBST(root.left, low, high);
        root.right = trimBST(root.right, low, high);
        return root;
    }
}

复杂度分析

• 时间复杂度：$O(n)$，其中 $n$ 是二叉搜索树的结点数。每个结点最多被访问一次。

• 空间复杂度：$O(n)$，其中 $n$ 是二叉搜索树的结点数。空间复杂度主要是栈空间，取决于二叉搜索树的高度，最坏情况下二叉搜索树的高度是 $O(n)$。

解法二

思路和算法

使用迭代实现修剪二叉搜索树，需要首先定位到修剪后的二叉搜索树的根结点，如果新的根结点不为空，则对新的根结点的左子树和右子树执行修剪操作。

从根结点开始遍历。如果当前结点值不在给定的边界范围内，则不保留当前结点，根据当前结点值和边界范围的大小关系决定移动到左子结点或者右子结点，直到当前结点变为空或者当前结点值在给定的边界范围内时，定位到新的根结点。

• 如果新的根结点为空，则返回空树。

• 如果新的根结点不为空，则对新的根结点的左子树和右子树执行修剪操作。

修剪左子树时，由于左子树中的所有结点值都小于根结点值，因此需要判断左子树中的每个结点值是否小于下界。将父结点初始化为新的根结点，将当前结点初始化为左子树的根结点。修剪操作如下。

• 如果当前结点值小于下界，则不保留当前结点和当前结点的左子树，因此将当前结点移动到当前结点的右子结点，将父结点的左子结点设为更新后的当前结点。

• 如果当前结点值大于等于下界，则保留当前结点，因此将父结点移动到当前结点，将当前结点移动到当前结点的左子结点。

修剪右子树时，由于右子树中的所有结点值都大于根结点值，因此需要判断右子树中的每个结点值是否大于上界。将父结点初始化为新的根结点，将当前结点初始化为右子树的根结点。修剪操作如下。

• 如果当前结点值大于上界，则不保留当前结点和当前结点的右子树，因此将当前结点移动到当前结点的左子结点，将父结点的右子结点设为更新后的当前结点。

• 如果当前结点值小于等于上界，则保留当前结点，因此将父结点移动到当前结点，将当前结点移动到当前结点的右子结点。

代码

class Solution {
    public TreeNode trimBST(TreeNode root, int low, int high) {
        while (root != null) {
            if (root.val < low) {
                root = root.right;
            } else if (root.val > high) {
                root = root.left;
            } else {
                break;
            }
        }
        if (root == null) {
            return null;
        }
        TreeNode parentLeft = root, childLeft = root.left;
        while (childLeft != null) {
            if (childLeft.val < low) {
                childLeft = childLeft.right;
                parentLeft.left = childLeft;
            } else {
                parentLeft = childLeft;
                childLeft = childLeft.left;
            }
        }
        TreeNode parentRight = root, childRight = root.right;
        while (childRight != null) {
            if (childRight.val > high) {
                childRight = childRight.left;
                parentRight.right = childRight;
            } else {
                parentRight = childRight;
                childRight = childRight.right;
            }
        }
        return root;
    }
}

复杂度分析

• 时间复杂度：$O(n)$，其中 $n$ 是二叉搜索树的结点数。每个结点最多被访问一次。

• 空间复杂度：$O(1)$。