【Lintcode】578. Lowest Common Ancestor III

题目地址：

https://www.lintcode.com/problem/lowest-common-ancestor-iii/description

给定一棵二叉树和节点$A$和$B$，求其最近公共祖先。如果不存在则返回null。注意，题目不保证$A$和$B$一定存在在二叉树中。

思路是，先判断$A$和$B$在不在二叉树中，如果有一个不在，直接返回null。否则开始考虑下一步操作。如果树为空，则返回null，否则，如果$A$正好是树根，那显然最近公共祖先就是$A$，对于$B$亦然。否则，递归调用，在左子树和右子树分别寻找公共祖先。具体请看代码：

public class Solution {
    /*
     * @param root: The root of the binary tree.
     * @param A: A TreeNode
     * @param B: A TreeNode
     * @return: Return the LCA of the two nodes.
     */
    public TreeNode lowestCommonAncestor3(TreeNode root, TreeNode A, TreeNode B) {
        // write your code here
        // 如果A或B其中一个不在树里，直接返回null
        if (!exists(root, A) || !exists(root, B)) {
            return null;
        }
        
        return lca(root, A, B);
    }
    
    // 这个函数的作用是，如果在以root为根的子树里能找到A和B的最近公共祖先，则返回之；
    // 如果该子树只含A或B其中一个节点，则返回那个节点；
    // 如果该子树既不包含A也不包含B，则返回null
    private TreeNode lca(TreeNode root, TreeNode A, TreeNode B) {
    	// 如果是空树，则返回null
        if (root == null) {
            return null;
        }
    	
    	// 如果A恰好是树根，那么显然最近公共祖先就是树根，则返回A；对于B亦然
        if (root == A) {
            return A;
        } else if (root == B) {
            return B;
        }
        
    	// 递归调用，在左子树和右子树中分别寻找
        TreeNode l = lca(root.left, A, B);
        TreeNode r = lca(root.right, A, B);
    	// 如果l为null，说明A和B都在右子树中，说明最近公共祖先也必然在右子树中，那么返回r；
    	// 如果r为null亦然；
    	// 如果都不为null，说明A和B在root两边，则返回root。
        if (l == null) {
            return r;
        } else if (r == null) {
            return l;
        } else {
            return root;
        }
    }
    
    private boolean exists(TreeNode root, TreeNode A) {
        if (root == null) {
            return false;
        }
        
        if (root == A) {
            return true;
        } else {
            return exists(root.left, A) || exists(root.right, A);
        }
    }
}

class TreeNode {
    int val;
    TreeNode left, right;
    TreeNode(int val) {
        this.val = val;
    }
}

时间复杂度$O(n)$，空间$O(h)$。