算法_二叉树_修剪二叉搜索树

修剪二叉搜索树

leetcode链接

1.两种解法

递归法

思路：这道题我们最直接的想法就是，对于一个二叉搜索树而言

如果某节点的值比low还要小，那么该节点的左子树就可以直接剪掉，因为左子树所有的节点肯定都比low小；
如果某节点的值比right要大，那么该节点的右子树也可以直接剪掉，因为右子树所有的节点肯定都比high大；
如果是在[low,high]之间，那么我们就要对该节点的左孩子和右孩子重复和上面一样的操作。

有了基本思路，我们可以写出以下代码：

def trimBST(self, root: Optional[TreeNode], low: int, high: int) -> Optional[TreeNode]:
	if not root:
	    return None
	
	if root.val<=low:
	    root.left = None # 剪掉左子树
	    self.trimBST(root.right,low,high) # 只递归右子树即可
	elif root.val>=high:
	    root.right = None # 剪掉右子树
	    self.trimBST(root.left,low,high) # 只递归左子树即可
	else:
	    self.trimBST(root.left,low,high) # 递归左子树
	    self.trimBST(root.right,low,high) # 递归右子树
	
	return root

下面我们来思考一些细节问题：

对于root.val<=low而言：

1. 当root.val==low时，那么我们执行上面的操作没有问题。
2. 当root.val<low时，那么我们既要剪掉左子树，也要把该节点剪掉，那么我们怎样才能剪掉该节点（node）呢？我们需要用一个pre节点记录该节点的父节点，这样让pre.left = node.right，即可把该节点和该节点的左子树全部剪掉

对于root.val>=high是同样的道理。只不过是当root.val>high时，执行pre.right = node.left。

例如：对下面的树只保留[1,3]之间的节点，那么我们删除0，就要把0的父节点3连接0的右孩子2，即pre.left = node.right。

看似现在没问题了，但是我们忽略了一个问题，就是每个节点都存在pre吗，即每个节点都有父节点吗？其实根节点就不存在父节点，所以当我们操作的时候要把父节点单独拿出来进行操作，但是这样就使得代码显得不简洁，所以不如我们先把根节点移动到一个一定在[low,high]内的节点的位置，这样根节点就不可能会有root.val<low或者root.val>high的情况了，所以就更谈不上pre了。

所以代码修改之后如下：

def trimBST(root, low, high):

    if not root:
        return None

	# 找到一个在[low,high]之间的节点作为root
    while root and (root.val<low or root.val>high):
        if root.val < low:
            root = root.right
        else:
            root = root.left

	# 因为在leetcode上提交最好不要修改他给的函数的参数，所以我就又写了一个函数，这个函数的参数带pre
    def traversal(root,pre,low,high):

        if not root:
            return None
		
		# 如果root.val <= low,存在两种情况，小于和等于
        if root.val <= low:
            root.left = None # 无论小于或者等于，都要剪掉左子树
            if root.val < low: # 如果小于
                pre.left = root.right # 则也要剪掉当前节点
                traversal(root.right,pre,low,high) # 然后遍历右子树，pre不变，因为把该节点删了，所以父节点依旧是之前的节点
            else: # 如果等于
                traversal(root.right,root,low,high) # 不用剪掉，直接遍历右子树，注意此时的pre要变为root，因为root.right的父节点是root

		# 如果root.val >= high,存在两种情况，大于和等于
        elif root.val >= high:
            root.right = None # 无论大于或者等于，都要剪掉右子树
            if root.val > high: # 如果大于
                pre.right = root.left # 就要把当前节点也剪掉
                traversal(root.left,pre,low,high) # 遍历左子树
            else: # 如果等于
                traversal(root.left,root,low,high) # 直接遍历左子树，pre变为root
        
        # 如果low < root.val < high
        else:
            traversal(root.left,root, low, high) # 遍历左子树
            traversal(root.right,root, low, high) # 遍历右子树

        return root

    root = traversal(root,None,low,high)

    return root

迭代法

思路：迭代法的思路和上面基本一致，我没写迭代法的python代码，可以参考下面的c++代码

TreeNode* trimBST(TreeNode* root, int L, int R) {
	if (!root) return nullptr;
	
	// 处理头结点，让root移动到[L, R] 范围内，注意是左闭右闭
	while (root != nullptr && (root->val < L || root->val > R)) {
	    if (root->val < L) root = root->right; // 小于L往右走
	    else root = root->left; // 大于R往左走
	}
	
	TreeNode *cur = root;
	// 此时root已经在[L, R] 范围内，处理左孩子元素小于L的情况
	while (cur != nullptr) {
	    while (cur->left && cur->left->val < L) {
	        cur->left = cur->left->right;
	    }
	    cur = cur->left;
	}
	cur = root;
	
	// 此时root已经在[L, R] 范围内，处理右孩子大于R的情况
	while (cur != nullptr) {
	    while (cur->right && cur->right->val > R) {
	        cur->right = cur->right->left;
	    }
	    cur = cur->right;
	}
	return root;
}

2.总结

算法

遇到根节点和其他的节点的的处理方法有一点区别的时候，可以先想想可不可以先遍历找到一个合适的根节点来避开有区别的地方。