Python long but easy to understand solution (beats 95%)


  • 0

    This should work even if p or q is not found in the tree, or p and q points to the same node.
    The code looks messy but while reading from top to bottom, it should be quite straight-forward recursion.

    Cracking the code interview uses a similar solution.

    class Solution(object):
        def lowestCommonAncestor(self, root, p, q):
            """
            :type root: TreeNode
            :type p: TreeNode
            :type q: TreeNode
            :rtype: TreeNode
            """
            x, pok, qok = self.find(root, p, q)
            if x is None:
                return None
            return x
    
        # find returns three values, x, pok, qok.
        # x returns TreeNode if LCA found, otherwise None.
        # pok returns True if p found under the node, otherwise False.
        # qok returns True if q found under the node, otherwise False.
        def find(self, parent, p, q):
            if not parent:
                return None, False, False
            if parent == p:
                if p == q:
                    return parent, True, True
                # parent is p, check q is found among left subtree or right subtree.
                _, _, qok1 = self.find(parent.right, p, q)
                if qok1:
                    return parent, True, True
                _, _, qok2 = self.find(parent.left, p, q)
                if qok2:
                    return parent, True, True
                return None, True, False
    
            if parent == q:
                if p == q:
                    return parent, True, True
                # parent is q, check p is found among left subtree or right subtree.    
                _, pok1, _ = self.find(parent.right, p, q)
                if pok1:
                    return parent, True, True
                _, pok2, _ = self.find(parent.left, p, q)
                if pok2:
                    return parent, True, True
                return None, False, True
    
            # check both right subtree and left subtree
            r, pok1, qok1 = self.find(parent.right, p, q)
            if r:
                return r, True, True
            l, pok2, qok2 = self.find(parent.left, p, q)
            if l:
                return l, True, True
            if (pok1 and qok2) or (
                    pok2 and qok1):  # if p and q found on each side, parent is LCA
                return parent, True, True
            return None, pok1 or pok2, qok1 or qok2
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.