python - trim bst

  • 0
    # 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 preorder(self, root, pre):
            if root:
                self.preorder(root.left, pre)
                self.preorder(root.right, pre)
        def find_next_greatest(self, pre, root_val):
            ret = 1
            for ret,_ in enumerate(pre):
                if pre[ret] > root_val:
                    return ret
            return ret+1
        def trim(self, pre, L, R):
            if not pre or L > R:
                return None
            root_val, node = pre[0], None
            ret = self.find_next_greatest(pre, root_val)
            if root_val < L:
                node = self.trim(pre[ret:], L,R)
            elif root_val > R:
                node = self.trim(pre[1:ret], L, R)
            elif L <= root_val <= R:    
                node = TreeNode(root_val)
                node.left = self.trim(pre[1:ret], L, root_val-1)
                node.right = self.trim(pre[ret:], root_val+1, R)
            return node
        def trimBST(self, root, L, R):
            :type root: TreeNode
            :type L: int
            :type R: int
            :rtype: TreeNode
            # construct a new tree with numbers between L and R
            # instead of deleting nodes from the current tree
            pre = []
            self.preorder(root, pre)
            return self.trim(pre, L, R)

Log in to reply

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.