701. Insert into a Binary Search Tree

Information

  • Diffculty: Medium

  • Created: 2020-02-13 10:10:00

  • Last Motified: 2020-02-13 10:19:54

  • Tags: Tree

Intuition

First, use a binary search to find the largest node which is smaller than the inserting value. Then, use this value to replace the found node, and add its right sub-tree to the right of the inserted one.

Solution

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

class Solution:
    def insertIntoBST(self, root: TreeNode, val: int) -> TreeNode:
        new = TreeNode(val)
        if root is None:
            return new

        # Find the closest node to the inserting value which is smaller than it
        closest = None
        diff = float('inf')
        tmp = root
        while tmp:
            if tmp.val > val:
                tmp = tmp.left
            else:
                if diff > val - tmp.val:
                    diff = val - tmp.val
                    closest = tmp
                tmp = tmp.right

        if closest is None:
            # All nodes in the tree is greater than val
            new.right = root
            return new

        new.right = closest.right
        closest.right = new
        return root