[leetcode] 776. Split BST @ python

原题

https://leetcode.com/problems/split-bst/

解法

递归, base case是根节点为空时, 返回[None, None]. 然后使用分治法, 分情况讨论.

代码

# 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 splitBST(self, root, V):
        """
        :type root: TreeNode
        :type V: int
        :rtype: List[TreeNode]
        """
        if not root:
            return [None, None]
        if root.val == V:
            temp = root.right
            root.right = None
            return [root, temp]
        elif root.val > V:
            small, large = self.splitBST(root.left, V)
            root.left = large
            return [small, root]
        else:
            small, large = self.splitBST(root.right, V)
            root.right = small
            return [root, large]