Lowest Common Ancestor

LeetCode 236 · View on LeetCode

Recursive DFS: if the current node is p or q, return it. Recurse into both subtrees. If both sides return non-None, current node is the LCA. Otherwise return whichever side found something.

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

def lowest_common_ancestor(root: TreeNode, p: TreeNode, q: TreeNode) -> TreeNode:
    if not root or root == p or root == q:
        return root
    left = lowest_common_ancestor(root.left, p, q)
    right = lowest_common_ancestor(root.right, p, q)
    if left and right:
        return root
    return left or right

if __name__ == '__main__':
    #       3
    #      / \
    #     5   1
    #    / \
    #   6   2
    n6 = TreeNode(6)
    n2 = TreeNode(2)
    n5 = TreeNode(5, n6, n2)
    n1 = TreeNode(1)
    root = TreeNode(3, n5, n1)

    print(lowest_common_ancestor(root, n5, n1).val)  # 3
    print(lowest_common_ancestor(root, n5, n2).val)  # 5