236. Lowest Common Ancestor of a Binary Tree - cocoder39/coco39_LC GitHub Wiki

class Solution:
    def lowestCommonAncestor(self, root: 'TreeNode', p: 'TreeNode', q: 'TreeNode') -> 'TreeNode':
        return self.findNode(root, p, q)

    # traverse from root and find first node of p or q
    def findNode(self, root: 'TreeNode', p: 'TreeNode', q: 'TreeNode'):
        if not root:
            return root
        
        if root == p or root == q: # return first node that is p/q
            return root
        
        left = self.findNode(root.left, p, q)
        right = self.findNode(root.right, p, q)
        if left and right: # p and q are in separate sides
            return root
        if left: # both p and q are in left side
            return left # left is first node of p/q
        return right

class Solution:
    def lowestCommonAncestor(self, root: 'TreeNode', p: 'TreeNode', q: 'TreeNode') -> 'TreeNode':
        if not root:
            return None
        
        # root is candidate of LCA
        # further searching down is unnecessary
        if 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 if left else right

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

class Solution {
public:
    TreeNode* lowestCommonAncestor(TreeNode* root, TreeNode* p, TreeNode* q) {
        if (find(root, p, q) != 2) { 
            return nullptr;
        }
        return helper(root, p, q);
    }
private:
    int find(TreeNode* root, TreeNode* p, TreeNode* q) {
        int res = 0;
        stack<TreeNode*> s;
        TreeNode* cur = root;
        while (cur || ! s.empty()) {
            if (cur) {
                res += (cur == p || cur == q);
                s.push(cur);
                cur = cur->left;
            } else {
                TreeNode* node = s.top();
                s.pop();
                cur = node->right;
            }
        }
        return res;
    }
    
    TreeNode* helper(TreeNode* root, TreeNode* p, TreeNode* q) {
        if (! root || root == p || root == q) { //find p or find q or find neither
            return root;
        }
        TreeNode* left = helper(root->left, p, q);
        TreeNode* right = helper(root->right, p, q);
        if (left && right) {
            return root;
        } else if (left) {
            return left;
        } else {
            return right;
        }
    }
};