LeetCode 1644. Lowest Common Ancestor of a Binary Tree II

Question

Given the root of a binary tree, return the lowest common ancestor (LCA) of two given nodes, p and q. If either node p or q does not exist in the tree, return null. All values of the nodes in the tree are unique.

According to the definition of LCA on Wikipedia: “The lowest common ancestor of two nodes p and q in a binary tree T is the lowest node that has both p and q as descendants (where we allow a node to be a descendant of itself)”. A descendant of a node x is a node y that is on the path from node x to some leaf node.

Example 1:

Input: root = [3,5,1,6,2,0,8,null,null,7,4], p = 5, q = 1
Output: 3
Explanation: The LCA of nodes 5 and 1 is 3.

Example 2:

Input: root = [3,5,1,6,2,0,8,null,null,7,4], p = 5, q = 4
Output: 5
Explanation: The LCA of nodes 5 and 4 is 5. A node can be a descendant of itself according to the definition of LCA.

Example 3:

Input: root = [3,5,1,6,2,0,8,null,null,7,4], p = 5, q = 10
Output: null
Explanation: Node 10 does not exist in the tree, so return null. 

Constraints:

• The number of nodes in the tree is in the range [1, 104].
• -109 <= Node.val <= 109
• All Node.val are unique.
• p != q

Source: https://leetcode.com/problems/lowest-common-ancestor-of-a-binary-tree-ii/

Solution

The solution in LeetCode 236. Lowest Common Ancestor of a Binary Tree cannot solve this question because p or q may not exist in the tree. The following solution separates the information for existence and the information for result (lowest common ancestor reference), so that fix that flaw.

static class TreeNode {
    int val;
    TreeNode left;
    TreeNode right;

    TreeNode(int x) {
        val = x;
    }
}

private TreeNode lca = null;

// return the number of target nodes in the tree of root
// separate lca information and target occurrence information
private int dfs(TreeNode root, TreeNode p, TreeNode q) {
    if (root == null) {
        return 0;
    }

    int count = 0;
    if (root == p || root == q) {
        count++;
    }
    count += dfs(root.left, p, q);
    count += dfs(root.right, p, q);

    // only set lca once when encountering the lowest common ancestor
    if (count == 2 && lca == null) {
        lca = root;
    }
    return count;
}

public TreeNode lowestCommonAncestor(TreeNode root, TreeNode p, TreeNode q) {
    lca = null;
    dfs(root, p, q);
    return lca;
}