[leetcode] 669. Trim a Binary Search Tree @ python

原题

Given a binary search tree and the lowest and highest boundaries as L and R, trim the tree so that all its elements lies in [L, R] (R >= L). You might need to change the root of the tree, so the result should return the new root of the trimmed binary search tree.

Example 1:
Input:
1
/
0 2

L = 1
R = 2

Output:
1

2
Example 2:
Input:
3
/
0 4

2
/
1

L = 1
R = 3

Output:
3
/
2
/
1

解法

递归, 使用分治法. Base case是当root为空时, 返回None. 分三种情况: 1) root的值比L小, 那么去掉root和它的左子树, 修剪右子树 2) root的值比R大, 那么去掉root和它的右子树, 修剪左子树 3) root的值在L和R之间, 那么保留root, 修剪左右子树.
Time: O(n) n为节点数
Space: O(1)

代码

# Definition for a binary tree node.
# class TreeNode(object):
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution(object):
    def trimBST(self, root, L, R):
        """
        :type root: TreeNode
        :type L: int
        :type R: int
        :rtype: TreeNode
        """
        # base case
        if not root:
            return None
        if root.val < L:
            # remove the root and its left subtree
            root = self.trimBST(root.right, L, R)
        elif root.val > R:
            # remove the root and its right subtree
            root = self.trimBST(root.left, L, R)
        else:
            # keep root and trim the left and right subtree
            root.left = self.trimBST(root.left, L, R)
            root.right = self.trimBST(root.right, L, R)
        return root