Leetcode

Lowest Common Ancestor of a Binary Tree IV

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

C++

class Solution {
 public:
  TreeNode* lowestCommonAncestor(TreeNode* root, vector<TreeNode*>& nodes) {
    unordered_set<TreeNode*> nodesSet{begin(nodes), end(nodes)};
    return lca(root, nodesSet);
  }

 private:
  TreeNode* lca(TreeNode* root, unordered_set<TreeNode*>& nodesSet) {
    if (!root)
      return nullptr;
    if (nodesSet.count(root))
      return root;

    TreeNode* l = lca(root->left, nodesSet);
    TreeNode* r = lca(root->right, nodesSet);

    if (l && r)
      return root;
    return l ? l : r;
  }
};

JAVA

class Solution {
  public TreeNode lowestCommonAncestor(TreeNode root, TreeNode[] nodes) {
    Set<TreeNode> nodesSet = new HashSet<>(Arrays.asList(nodes));
    return lca(root, nodesSet);
  }

  private TreeNode lca(TreeNode root, Set<TreeNode> nodesSet) {
    if (root == null)
      return null;
    if (nodesSet.contains(root))
      return root;

    TreeNode l = lca(root.left, nodesSet);
    TreeNode r = lca(root.right, nodesSet);

    if (l != null && r != null)
      return root;
    return l == null ? r : l;
  }
}

Python

class Solution:
  def lowestCommonAncestor(self, root: 'TreeNode', nodes: 'List[TreeNode]') -> 'TreeNode':
    nodes = set(nodes)

    def lca(root: 'TreeNode') -> 'TreeNode':
      if not root:
        return None
      if root in nodes:
        return root

      l = lca(root.left)
      r = lca(root.right)

      if l and r:
        return root
      return l or r

    return lca(root)