LeetCode 236: Lowest Common Ancestor of a Binary Tree

Problem Restatement

We are given:

• The root of a binary tree
• Two nodes p and q

We need to return their lowest common ancestor.

The lowest common ancestor is the lowest node in the tree that has both p and q as descendants. A node can be a descendant of itself.

Unlike LeetCode 235, this tree is a normal binary tree, not a binary search tree. So we cannot use value ordering to decide whether to go left or right.

LeetCode states that p and q both exist in the tree, p != q, and all node values are unique. The number of nodes is in the range [2, 10^5].

Input and Output

Tree node shape:

class TreeNode:
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None

Function shape:

class Solution:
    def lowestCommonAncestor(
        self,
        root: "TreeNode",
        p: "TreeNode",
        q: "TreeNode",
    ) -> "TreeNode":
        ...

Examples

Example 1:

Input:
root = [3,5,1,6,2,0,8,null,null,7,4]
p = 5
q = 1

Output:
3

Tree:

        3
      /   \
     5     1
    / \   / \
   6   2 0   8
      / \
     7   4

Node 5 is in the left subtree of 3.

Node 1 is in the right subtree of 3.

So their lowest common ancestor is:

3

Example 2:

Input:
root = [3,5,1,6,2,0,8,null,null,7,4]
p = 5
q = 4

Output:
5

Node 5 is an ancestor of node 4.

Because a node can be a descendant of itself, the answer is:

5

First Thought: Store Root-to-Node Paths

A direct method is:

1. Find the path from root to p.
2. Find the path from root to q.
3. Compare both paths from the beginning.
4. The last equal node is the LCA.

This works, but it needs extra space for both paths and requires more bookkeeping.

We can solve it more directly with recursive DFS.

Key Insight

At any node, there are only a few important cases:

1. The current node is None.
2. The current node is p or q.
3. One target is found in the left subtree.
4. One target is found in the right subtree.
5. Targets are found on both sides.

The important case is when one target is found on the left and the other target is found on the right.

Then the current node is the lowest common ancestor.

Recursive Meaning

Define this function:

lowestCommonAncestor(root, p, q)

to return:

This return meaning lets recursion carry information upward.

Algorithm

For a node root:

1. If root is None, return None.
2. If root is p or q, return root.
3. Search the left subtree.
4. Search the right subtree.
5. If both sides return non-null, return root.
6. If only one side returns non-null, return that side.
7. If both sides return null, return None.

Why This Works

Suppose the left recursive call returns a node and the right recursive call returns a node.

That means one target was found in the left subtree and the other target was found in the right subtree.

The current node is the first node where the two search paths meet.

So the current node is the LCA.

If only the left side returns a node, then both targets are either in the left subtree, or only one target has been found so far. In both cases, the correct information to pass upward is the left result.

The same logic applies to the right side.

If the current node itself is p or q, we return it immediately. This correctly handles the case where one target is an ancestor of the other.

Correctness

We prove the algorithm is correct by considering each subtree.

For an empty subtree, the algorithm returns None, which is correct because no target node exists there.

For a non-empty subtree rooted at root, if root is p or q, returning root is correct because a target node was found. If the other target appears below it, this node will become the ancestor returned upward.

Otherwise, the algorithm recursively searches the left and right subtrees.

If both recursive calls return non-null values, then one target was found on each side. Since the left and right subtrees are disjoint, the current node is the lowest node that contains both targets. So returning root is correct.

If exactly one recursive call returns a non-null value, then all found target information lies on that side. Returning that non-null value preserves the only possible LCA or target found so far.

If both calls return None, neither target is in this subtree, so returning None is correct.

Because p and q are guaranteed to exist in the tree, the final result returned from the original root is the lowest common ancestor.

Complexity

Here, n is the number of nodes and h is the height of the tree.

For a balanced tree, h = O(log n).

For a skewed tree, h = O(n).

Implementation

class Solution:
    def lowestCommonAncestor(
        self,
        root: "TreeNode",
        p: "TreeNode",
        q: "TreeNode",
    ) -> "TreeNode":
        if root is None:
            return None

        if root == p or root == q:
            return root

        left = self.lowestCommonAncestor(root.left, p, q)
        right = self.lowestCommonAncestor(root.right, p, q)

        if left and right:
            return root

        if left:
            return left

        return right

Code Explanation

The empty subtree case is:

if root is None:
    return None

Then we check whether the current node is one of the targets:

if root == p or root == q:
    return root

This handles ancestor cases, such as p = 5 and q = 4, where node 5 itself is the answer.

Next, search both subtrees:

left = self.lowestCommonAncestor(root.left, p, q)
right = self.lowestCommonAncestor(root.right, p, q)

If both sides found something:

if left and right:
    return root

then the current node joins the two paths.

If only one side found something:

if left:
    return left

return right

we pass that result upward.

Testing

from collections import deque

class TreeNode:
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None

def build_tree(values):
    if not values:
        return None

    root = TreeNode(values[0])
    q = deque([root])
    i = 1

    while q and i < len(values):
        node = q.popleft()

        if i < len(values) and values[i] is not None:
            node.left = TreeNode(values[i])
            q.append(node.left)
        i += 1

        if i < len(values) and values[i] is not None:
            node.right = TreeNode(values[i])
            q.append(node.right)
        i += 1

    return root

def find_node(root, value):
    if root is None:
        return None

    if root.val == value:
        return root

    return find_node(root.left, value) or find_node(root.right, value)

def run_tests():
    s = Solution()

    root = build_tree([3, 5, 1, 6, 2, 0, 8, None, None, 7, 4])

    p = find_node(root, 5)
    q = find_node(root, 1)
    assert s.lowestCommonAncestor(root, p, q).val == 3

    p = find_node(root, 5)
    q = find_node(root, 4)
    assert s.lowestCommonAncestor(root, p, q).val == 5

    p = find_node(root, 7)
    q = find_node(root, 4)
    assert s.lowestCommonAncestor(root, p, q).val == 2

    p = find_node(root, 6)
    q = find_node(root, 4)
    assert s.lowestCommonAncestor(root, p, q).val == 5

    root = build_tree([1, 2])
    p = find_node(root, 1)
    q = find_node(root, 2)
    assert s.lowestCommonAncestor(root, p, q).val == 1

    print("all tests passed")

run_tests()