leetcode 669 修剪二叉搜索树

题目

给你二叉搜索树的根节点 root ，同时给定最小边界low 和最大边界 high。通过修剪二叉搜索树，使得所有节点的值在[low, high]中。修剪树 不应该 改变保留在树中的元素的相对结构 (即，如果没有被移除，原有的父代子代关系都应当保留)。 可以证明，存在 唯一的答案 。

所以结果应当返回修剪好的二叉搜索树的新的根节点。注意，根节点可能会根据给定的边界发生改变。

示例

输入：root = [1,0,2], low = 1, high = 2
输出：[1,null,2]

解析

首先可以想到递归，由于是二叉搜索树，比如递归到某个节点的时候，值小于最小值，这时还需要去判断该节点的右子树是否小于该值（二叉搜索树的特性）；大于同理：

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
func trimBST(root *TreeNode, low int, high int) *TreeNode {
    if root == nil {
        return nil
    }
    if root.Val < low {
        return trimBST(root.Right, low, high)
    }
    if root.Val > high {
        return trimBST(root.Left, low, high)
    }
    root.Left = trimBST(root.Left, low, high)
    root.Right = trimBST(root.Right, low, high)
    return root
}

迭代法有那么一点点不太好想：

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
func trimBST(root *TreeNode, low int, high int) *TreeNode {
    if root == nil {
        return nil
    }
    // 处理 root，让 root 移动到[low, high] 范围内
    for root != nil && (root.Val < low || root.Val > high) {
        if root.Val < low {
            root = root.Right
        } else {
            root = root.Left
        }
    }
    //处理左孩子元素小于 low 的情况
    cur := root 
    for cur != nil {
        for cur.Left != nil && cur.Left.Val < low {
            cur.Left = cur.Left.Right
        }
        cur = cur.Left
    }
    //处理右孩子大于 high 的情况
    cur = root
    for cur != nil {
        for cur.Right != nil && cur.Right.Val > high {
            cur.Right = cur.Right.Left
        }
        cur = cur.Right
    }
    return root
}