236. Lowest Common Ancestor of a Binary Tree

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

Given a binary tree, find the lowest common ancestor (LCA) of two given nodes in the tree.

According to the definition of LCA on Wikipedia: “The lowest common ancestor is defined between two nodes v and w as the lowest node in T that has both v and w as descendants (where we allow a node to be a descendant of itself).”

Thoughts

找二叉树上两个结点的最底层的共同祖先。分治思想：当前结点为lca和lca分别在左右子树，以及lca不在此子树上。让lca()遇到p/q时直接返回p/q，此时当左右都不为null时当前结点就是LCA，否则定有一个子树返回空。

Code

/*
 * @lc app=leetcode id=236 lang=cpp
 *
 * [236] Lowest Common Ancestor of a Binary Tree
 */

// @lc code=start
/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    TreeNode* lowestCommonAncestor(TreeNode* root, TreeNode* p, TreeNode* q) {
        if (root == nullptr || root == p || root == q) return root;
        auto left = lowestCommonAncestor(root->left, p, q); 
        auto right = lowestCommonAncestor(root->right, p, q);
        if (left != nullptr && right != nullptr) return root;
        return left == nullptr ? right : left;
    }
};
// @lc code=end

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);
        return left == null ? right : right == null ? left : root;
    }
}

Analysis

只有一次分治，n个结点每个结点O(1)，总时间O(n).

Ver.2

ans stack存储每个节点是否是p和q的ancestor，巧妙利用pre结点已经包含了两种backtracking的结果让它和当前结果做OR。由于pre还可能是上步push_left(cur->right)中的cur，此时的cur一定是叶子，因此对push_left后这个特殊pre做下清空。

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    TreeNode* lowestCommonAncestor(TreeNode* root, TreeNode* p, TreeNode* q) {
        if (p == q) return p;
        stack<TreeNode*> s;
        stack<pair<bool, bool>> ans;
        const auto push_left = [&](TreeNode *cur) {
            while (cur != NULL) {
                s.push(cur);
                if (cur == p) ans.push({true, false});
                else if (cur == q) ans.push({false, true});
                else ans.push({false, false});
                cur = cur->left;
            }
        };
        TreeNode *pre;
        pair<bool, bool> pre_a;
        push_left(root);
        while (!s.empty()) {
            const auto cur = s.top();
            auto &a = ans.top();
            a.first |= pre_a.first;
            a.second |= pre_a.second;
            if (a.first && a.second) return cur;
            if (cur->right != NULL && cur->right != pre) {
                push_left(cur->right);
                pre_a = {false, false};
                continue;
            }
            pre = cur;
            pre_a = a;
            s.pop();
            ans.pop();
        }
        
        return root;
    }
};