DSA Lowest Common Ancestor in a BST

The problem

The lowest common ancestor of two nodes is the deepest node whose subtree contains both of them. Lowest means closest to the two nodes. A node counts as an ancestor of itself, so a target inside another target's subtree has that other target as its answer.

The tree in this problem is a binary search tree. Smaller values sit on the left and bigger values sit on the right at every node. That ordering is what makes a short walk possible.

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

def level_order(root):
    if root is None:
        return []
    out = []
    queue = [root]
    while queue:
        node = queue.pop(0)
        if node is None:
            out.append(None)
        else:
            out.append(node.value)
            queue.append(node.left)
            queue.append(node.right)
    while out and out[-1] is None:
        out.pop()
    return out

def insert(root, value):
    if root is None:
        return TreeNode(value)
    if value < root.value:
        root.left = insert(root.left, value)
    else:
        root.right = insert(root.right, value)
    return root

def build_bst(values):
    tree = None
    for v in values:
        tree = insert(tree, v)
    return tree

tree = build_bst([8, 3, 10, 1, 6, 14, 4, 7])
print("level order:", level_order(tree))
print("LCA of 1 and 7: expected 3")
print("LCA of 4 and 7: expected 6")
print("LCA of 1 and 14: expected 8")
print("LCA of 6 and 7: expected 6")
print("LCA of 1 and 4: expected 3")

The tree holds eight values. Its root is 8, with 3 and 10 below it. The five pairs above cover the shapes that matter. One pair splits at the root, one pair splits deep, and one pair has its answer at one of the targets itself.

Intuition

The first idea checks only the root. If the targets sit on different sides of the root, the root is the answer. If they sit on the same side, a root-only check has nothing more to say, and a guess of the root is wrong.

The pair 1 and 7 shows the gap. Both values are smaller than 8, so both live in the left subtree. Their common ancestor sits inside that subtree, deeper than the root.

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

def level_order(root):
    if root is None:
        return []
    out = []
    queue = [root]
    while queue:
        node = queue.pop(0)
        if node is None:
            out.append(None)
        else:
            out.append(node.value)
            queue.append(node.left)
            queue.append(node.right)
    while out and out[-1] is None:
        out.pop()
    return out

def insert(root, value):
    if root is None:
        return TreeNode(value)
    if value < root.value:
        root.left = insert(root.left, value)
    else:
        root.right = insert(root.right, value)
    return root

def build_bst(values):
    tree = None
    for v in values:
        tree = insert(tree, v)
    return tree

def root_only(root, a, b):
    if (a < root.value and b > root.value) or (a > root.value and b < root.value):
        return root.value
    if a == root.value or b == root.value:
        return root.value
    return root.value

tree = build_bst([8, 3, 10, 1, 6, 14, 4, 7])
print("pair 1 and 7 both smaller than root 8:", 1 < 8 and 7 < 8)
print("root-only answer for 1 and 7:", root_only(tree, 1, 7))
print("expected answer: 3")

The root-only answer is 8 and the true answer is 3. The fix is to repeat the same side test at each node on the way down. The walk keeps moving while both targets sit on one side, and it stops at the first split.

Brute force

A general route ignores the ordering completely. It searches the whole tree and records the parent of every node. It then collects every ancestor of the first target and walks up from the second target until it meets one of them.

The route works on any binary tree. Its cost is the full search. Every node is visited even when the two targets sit near the root.

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

def level_order(root):
    if root is None:
        return []
    out = []
    queue = [root]
    while queue:
        node = queue.pop(0)
        if node is None:
            out.append(None)
        else:
            out.append(node.value)
            queue.append(node.left)
            queue.append(node.right)
    while out and out[-1] is None:
        out.pop()
    return out

def insert(root, value):
    if root is None:
        return TreeNode(value)
    if value < root.value:
        root.left = insert(root.left, value)
    else:
        root.right = insert(root.right, value)
    return root

def build_bst(values):
    tree = None
    for v in values:
        tree = insert(tree, v)
    return tree

visits = 0

def brute_lca(root, a, b):
    global visits
    visits = 0
    parent = {root.value: None}
    order = [root]
    while order:
        node = order.pop(0)
        visits += 1
        for child in (node.left, node.right):
            if child is not None:
                parent[child.value] = node.value
                order.append(child)
    ancestors = set()
    x = a
    while x is not None:
        ancestors.add(x)
        x = parent[x]
    y = b
    while y not in ancestors:
        y = parent[y]
    return y

tree = build_bst([8, 3, 10, 1, 6, 14, 4, 7])
for a, b in [(1, 7), (4, 7), (1, 14), (6, 7)]:
    print("LCA of", a, "and", b, ":", brute_lca(tree, a, b), "visits:", visits)

Every answer is right and every search visits all eight nodes. The parent map also costs one entry per node. The BST ordering can cut both costs to a single downward path.

The optimal solution: one walk down

The key idea is one sentence. Walk down from the root, go left while both targets are smaller and right while both are bigger, and stop at the first node where they split or where one of them equals the node.

That stopping node is the answer. Everything above it holds both targets on the same side, so it is a shallower common ancestor. Everything below it can hold at most one of them.

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

def level_order(root):
    if root is None:
        return []
    out = []
    queue = [root]
    while queue:
        node = queue.pop(0)
        if node is None:
            out.append(None)
        else:
            out.append(node.value)
            queue.append(node.left)
            queue.append(node.right)
    while out and out[-1] is None:
        out.pop()
    return out

def insert(root, value):
    if root is None:
        return TreeNode(value)
    if value < root.value:
        root.left = insert(root.left, value)
    else:
        root.right = insert(root.right, value)
    return root

def build_bst(values):
    tree = None
    for v in values:
        tree = insert(tree, v)
    return tree

def lca(root, a, b):
    node = root
    while node is not None:
        if a < node.value and b < node.value:
            node = node.left
        elif a > node.value and b > node.value:
            node = node.right
        else:
            return node.value
    return None

tree = build_bst([8, 3, 10, 1, 6, 14, 4, 7])
for a, b in [(1, 7), (4, 7), (1, 14), (6, 7), (1, 4)]:
    print("LCA of", a, "and", b, ":", lca(tree, a, b))

All five cases from the first section pass. The walk never visits a subtree that cannot hold the answer. For the pair 1 and 14 it stops at once, because the targets split at the root itself.

Walkthrough

Follow the pair 4 and 7 through the tree. The root is 8. Both targets are smaller than 8, so the walk goes left before any answer is possible.

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

def level_order(root):
    if root is None:
        return []
    out = []
    queue = [root]
    while queue:
        node = queue.pop(0)
        if node is None:
            out.append(None)
        else:
            out.append(node.value)
            queue.append(node.left)
            queue.append(node.right)
    while out and out[-1] is None:
        out.pop()
    return out

def insert(root, value):
    if root is None:
        return TreeNode(value)
    if value < root.value:
        root.left = insert(root.left, value)
    else:
        root.right = insert(root.right, value)
    return root

def build_bst(values):
    tree = None
    for v in values:
        tree = insert(tree, v)
    return tree

def traced(root, a, b):
    node = root
    while node is not None:
        if a < node.value and b < node.value:
            print("at", node.value, ": both targets smaller, go left")
            node = node.left
        elif a > node.value and b > node.value:
            print("at", node.value, ": both targets bigger, go right")
            node = node.right
        else:
            print("at", node.value, ": targets split or match, answer is", node.value)
            return node.value
    return None

tree = build_bst([8, 3, 10, 1, 6, 14, 4, 7])
print("answer:", traced(tree, 4, 7))

At 8 the walk goes left. At 3 both targets are bigger, so the walk goes right. At 6 the targets split, with 4 on the left and 7 on the right. No deeper node can hold both, so 6 is the lowest common ancestor.

Complexity

Time for the walk is O(h) for a tree of height h. The walk follows one path from the root, so it makes one comparison step per level.

Space is O(1). The walk holds one current node and nothing else. The brute force pays O(n) time and O(n) space for its parent map on a tree of n nodes.

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

def level_order(root):
    if root is None:
        return []
    out = []
    queue = [root]
    while queue:
        node = queue.pop(0)
        if node is None:
            out.append(None)
        else:
            out.append(node.value)
            queue.append(node.left)
            queue.append(node.right)
    while out and out[-1] is None:
        out.pop()
    return out

def insert(root, value):
    if root is None:
        return TreeNode(value)
    if value < root.value:
        root.left = insert(root.left, value)
    else:
        root.right = insert(root.right, value)
    return root

def build_bst(values):
    tree = None
    for v in values:
        tree = insert(tree, v)
    return tree

steps = 0

def lca_counted(root, a, b):
    global steps
    steps = 0
    node = root
    while node is not None:
        steps += 1
        if a < node.value and b < node.value:
            node = node.left
        elif a > node.value and b > node.value:
            node = node.right
        else:
            return node.value
    return None

def balanced(sorted_vals):
    if not sorted_vals:
        return None
    mid = len(sorted_vals) // 2
    node = TreeNode(sorted_vals[mid])
    node.left = balanced(sorted_vals[:mid])
    node.right = balanced(sorted_vals[mid + 1:])
    return node

for n in [7, 15, 31]:
    tree = balanced(list(range(1, n + 1)))
    ans = lca_counted(tree, 1, n)
    print("nodes:", n, "answer:", ans, "walk steps:", steps, "full search visits:", n)

The node count doubles and then doubles again. The walk grows by one step each time, because a balanced tree gains one level. The full search grows with every node added, which is the gap the ordering buys.

Edge cases and pitfalls

The answer can be the root itself. That happens when the targets split at the root or when one target is the root. The walk handles it on its first step, with no special case.

The answer can also be one of the targets. The split test uses strict comparisons, so a node equal to a target stops the walk. Both targets must be present in the tree. The walk never checks the rest of the tree, so a missing target gives a meaningless answer instead of an error.

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

def level_order(root):
    if root is None:
        return []
    out = []
    queue = [root]
    while queue:
        node = queue.pop(0)
        if node is None:
            out.append(None)
        else:
            out.append(node.value)
            queue.append(node.left)
            queue.append(node.right)
    while out and out[-1] is None:
        out.pop()
    return out

def insert(root, value):
    if root is None:
        return TreeNode(value)
    if value < root.value:
        root.left = insert(root.left, value)
    else:
        root.right = insert(root.right, value)
    return root

def build_bst(values):
    tree = None
    for v in values:
        tree = insert(tree, v)
    return tree

def lca(root, a, b):
    node = root
    while node is not None:
        if a < node.value and b < node.value:
            node = node.left
        elif a > node.value and b > node.value:
            node = node.right
        else:
            return node.value
    return None

tree = build_bst([8, 3, 10, 1, 6, 14, 4, 7])
single = TreeNode(9)
print("split at root, 1 and 14:", lca(tree, 1, 14))
print("one target is the root, 8 and 7:", lca(tree, 8, 7))
print("answer is a target, 6 and 7:", lca(tree, 6, 7))
print("same target twice, 6 and 6:", lca(tree, 6, 6))
print("both deep right, 10 and 14:", lca(tree, 10, 14))
print("single-node tree, 9 and 9:", lca(single, 9, 9))

The same-target case stops at that node itself. A node is its own ancestor under this problem's rule, so the pair 6 and 6 answers 6. Forgetting that rule is the most common way to send the walk one step too deep.