236. Lowest Common Ancestor of a Binary Tree

Approach 1: Recursive

• Time: $O(h)$
• Space: $O(h)$

C++JavaPython

class Solution {
 public:
  TreeNode* lowestCommonAncestor(TreeNode* root, TreeNode* p, TreeNode* q) {
    if (root == nullptr || root == p || root == q)
      return root;
    TreeNode* left = lowestCommonAncestor(root->left, p, q);
    TreeNode* right = lowestCommonAncestor(root->right, p, q);
    if (left != nullptr && right != nullptr)
      return root;
    return left == nullptr ? right : left;
  }
};

class Solution {
  public TreeNode lowestCommonAncestor(TreeNode root, TreeNode p, TreeNode q) {
    if (root == null || root == p || root == q)
      return root;
    TreeNode left = lowestCommonAncestor(root.left, p, q);
    TreeNode right = lowestCommonAncestor(root.right, p, q);
    if (left != null && right != null)
      return root;
    return left == null ? right : left;
  }
}

class Solution:
  def lowestCommonAncestor(self, root: 'TreeNode', p: 'TreeNode', q: 'TreeNode') -> 'TreeNode':
    if not root or 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
    return left or right

Approach 2: Iterative

• Time: $O(n)$
• Space: $O(n)$

C++JavaPython

class Solution {
 public:
  TreeNode* lowestCommonAncestor(TreeNode* root, TreeNode* p, TreeNode* q) {
    queue<TreeNode*> q_{{root}};
    unordered_map<TreeNode*, TreeNode*> parent{{root, nullptr}};
    unordered_set<TreeNode*> ancestors;  // p's ancestors

    // Iterate until we find both p and q.
    while (!parent.count(p) || !parent.count(q)) {
      root = q_.front(), q_.pop();
      if (root->left) {
        parent[root->left] = root;
        q_.push(root->left);
      }
      if (root->right) {
        parent[root->right] = root;
        q_.push(root->right);
      }
    }

    // Insert all the p's ancestors.
    while (p) {
      ancestors.insert(p);
      p = parent[p];  // `p` becomes nullptr in the end.
    }

    // Go up from q until we meet any of p's ancestors.
    while (!ancestors.count(q))
      q = parent[q];

    return q;
  }
};

class Solution {
  public TreeNode lowestCommonAncestor(TreeNode root, TreeNode p, TreeNode q) {
    Queue<TreeNode> q_ = new ArrayDeque<>(Arrays.asList(root));
    Map<TreeNode, TreeNode> parent = new HashMap<>();
    parent.put(root, null);
    Set<TreeNode> ancestors = new HashSet<>(); // p's ancestors

    // Iterate until we find both p and q.
    while (!parent.containsKey(p) || !parent.containsKey(q)) {
      root = q_.poll();
      if (root.left != null) {
        parent.put(root.left, root);
        q_.offer(root.left);
      }
      if (root.right != null) {
        parent.put(root.right, root);
        q_.offer(root.right);
      }
    }

    // Insert all the p's ancestors.
    while (p != null) {
      ancestors.add(p);
      p = parent.get(p); // `p` becomes null in the end.
    }

    // Go up from q until we meet any of p's ancestors.
    while (!ancestors.contains(q))
      q = parent.get(q);

    return q;
  }
}

class Solution:
  def lowestCommonAncestor(self, root: Optional[TreeNode], p: Optional[TreeNode], q: Optional[TreeNode]) -> Optional[TreeNode]:
    q_ = collections.deque([root])
    parent = {root: None}
    ancestors = set()  # p's ancestors

    # Iterate until we find both p and q.
    while p not in parent or q not in parent:
      root = q_.popleft()
      if root.left:
        parent[root.left] = root
        q_.append(root.left)
      if root.right:
        parent[root.right] = root
        q_.append(root.right)

    # Insert all the p's ancestors.
    while p:
      ancestors.add(p)
      p = parent[p]  # `p` becomes None in the end.

    # Go up from q until we meet any of p's ancestors.
    while q not in ancestors:
      q = parent[q]

    return q