LeetCode 700: Search in a Binary Search Tree

Problem Restatement

We are given the root of a binary search tree and an integer val.

We need to find the node whose value equals val.

If the node exists, return that node. Returning the node means returning the entire subtree rooted at that node.

If the node does not exist, return null. The official statement asks us to return the subtree rooted at the node whose value equals val, or null if no such node exists.

Input and Output

Example function shape:

def searchBST(root: Optional[TreeNode], val: int) -> Optional[TreeNode]:
    ...

Examples

Example 1:

root = [4, 2, 7, 1, 3]
val = 2

Tree:

      4
     / \
    2   7
   / \
  1   3

The node with value 2 exists.

Return the subtree:

    2
   / \
  1   3

Answer:

[2, 1, 3]

Example 2:

root = [4, 2, 7, 1, 3]
val = 5

There is no node with value 5.

Answer:

None

First Thought: Search Every Node

A general binary tree search can use DFS.

We could visit every node and check whether its value equals val.

This works, but it ignores the binary search tree property.

A BST gives us a better rule.

Key Insight

In a binary search tree:

Why?

Every value in the left subtree is smaller than the current node.

Every value in the right subtree is larger than the current node.

So at each node, one comparison tells us which half of the tree can still contain the answer.

Algorithm

Start at the root.

While the current node is not None:

1. If current.val == val, return current.
2. If val < current.val, move to current.left.
3. Otherwise, move to current.right.

If the loop ends, the value was not found.

Return None.

Walkthrough

Consider:

root = [4, 2, 7, 1, 3]
val = 2

Start at node 4.

Since:

2 < 4

move left.

Now we are at node 2.

Since:

2 == 2

return this node.

The returned subtree is:

    2
   / \
  1   3

Now consider:

val = 5

Start at node 4.

5 > 4

move right to node 7.

At node 7:

5 < 7

move left.

The left child is None.

So 5 does not exist in the tree.

Return None.

Correctness

At every node, the binary search tree property tells us that all values in the left subtree are smaller than the current node, and all values in the right subtree are larger than the current node.

If val is smaller than the current node’s value, it cannot appear in the right subtree, so searching only the left subtree is safe.

If val is larger than the current node’s value, it cannot appear in the left subtree, so searching only the right subtree is safe.

If the current node’s value equals val, the algorithm returns exactly the required subtree rooted at that node.

If the search reaches None, every possible subtree that could contain val has been eliminated. Therefore, the value is not in the tree.

So the algorithm returns the matching node exactly when it exists, and None otherwise.

Complexity

Let h be the height of the tree.

If the BST is balanced, h = O(log n).

If the BST is skewed, h = O(n).

Implementation

class Solution:
    def searchBST(
        self,
        root: Optional[TreeNode],
        val: int,
    ) -> Optional[TreeNode]:
        current = root

        while current is not None:
            if current.val == val:
                return current

            if val < current.val:
                current = current.left
            else:
                current = current.right

        return None

Code Explanation

We begin at the root:

current = root

As long as there is a node to inspect:

while current is not None:

we compare its value with the target.

If it matches:

if current.val == val:
    return current

we return the node itself.

If the target is smaller:

if val < current.val:
    current = current.left

we move left.

Otherwise:

current = current.right

we move right.

If the loop finishes:

return None

the value was not found.

Recursive Version

The same logic can be written recursively.

class Solution:
    def searchBST(
        self,
        root: Optional[TreeNode],
        val: int,
    ) -> Optional[TreeNode]:
        if root is None:
            return None

        if root.val == val:
            return root

        if val < root.val:
            return self.searchBST(root.left, val)

        return self.searchBST(root.right, val)

This version is shorter, but it uses recursion stack space.

Testing

def run_tests():
    # Helper tree:
    #
    #       4
    #      / \
    #     2   7
    #    / \
    #   1   3
    #
    # searchBST(root, 2) should return the node with value 2.
    # Its subtree is [2, 1, 3].
    #
    # searchBST(root, 5) should return None.
    #
    # searchBST(root, 4) should return the root.
    #
    # searchBST(root, 1) should return the leaf node with value 1.
    #
    # searchBST(root, 8) should return None.

    pass

Test meaning: