lca.py

# Definition for a binary tree node.
class TreeNode(object):
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None
    def __repr__(self):
        return "TreeNode({})".format(self.val)

# see https://www.geeksforgeeks.org/lowest-common-ancestor-binary-tree-set-1/ for tips

class Solution(object):
    def lowestCommonAncestor(self, root, p, q):
        """
        :type root: TreeNode
        :type p: TreeNode
        :type q: TreeNode
        :rtype: TreeNode
        """
        # return self.lowestCommonAncestorPath(root, p, q)
        return self.lowestCommonAncestorRecursive(root, p, q)
    
    def lowestCommonAncestorRecursive(self, root, p, q):
        if root == None:
            return None
        elif root == p and self.subtreeContainsNode(p, [q]):
            return root 
        elif root == q and self.subtreeContainsNode(q, [p]):
            return root
        else:
            lefts = self.subtreeContainsNode(root.left, [p, q])
            rights = self.subtreeContainsNode(root.right, [p, q])

            if lefts and rights:
                return root
            elif lefts:
                return self.lowestCommonAncestorRecursive(root.left, p, q)
            elif rights:
                return self.lowestCommonAncestorRecursive(root.right, p, q)

    def subtreeContainsNode(self, root, nodes):
        if root == None:
            return False
        elif root in nodes:
            return True
        else:
            return self.subtreeContainsNode(root.left, nodes) \
                or self.subtreeContainsNode(root.right, nodes)


    def lowestCommonAncestorPath(self, root, p, q):
        ps = self.pathToNode(root, p)
        qs = self.pathToNode(root, q)

        if not ps or not qs:
            return None

        i = 0
        while i < len(ps) and i < len(qs) and ps[i] == qs[i]:
            i += 1
        
        return ps[i-1]

    def pathToNode(self, root, node):
        if root == None:
            return []
        elif root == node:
            return [node]
        else:
            left = self.pathToNode(root.left, node)
            if left:
                return [root] + left

            right = self.pathToNode(root.right, node)
            if right:
                return [root] + right

            if not left and not right:
                return []

def listToTree(lst):
    if len(lst) == 0:
        return None
    elif len(lst) == 1:
        return TreeNode(lst[0])

    root = TreeNode(lst[0])
    lstnodes = [root]
    for i in range(1, len(lst)):
        if lst[i] == None:
            node = None
        else:
            node = TreeNode(lst[i])

        index = int(i/2)

        print(node, index)

        if i % 2 == 1:
            lstnodes[index].left = node
        else:
            lstnodes[index].right = node
        lstnodes.append(node)

    return root

s = Solution()

t = TreeNode(3)

t.left = TreeNode(5)
t.left.left = TreeNode(6)
t.left.right = TreeNode(2)
t.left.right.left = TreeNode(7)
t.left.right.right = TreeNode(4)

t.right = TreeNode(1)
t.right.left = TreeNode(0)
t.right.right = TreeNode(8)

# print(s.subtreeContainsNode(t, t))
# print(s.subtreeContainsNode(t, t.left))
# print(s.subtreeContainsNode(t, t.left.right.right))
# print(s.subtreeContainsNode(t, TreeNode(11)))

# print(s.pathToNode(t, t.left.right.right))
# print(s.pathToNode(t, TreeNode(11)))

print(s.lowestCommonAncestor(t, t.left, t.right))
print(s.lowestCommonAncestor(t, t.left, t.left.right.right))
print(s.lowestCommonAncestor(t, t.left, TreeNode(11)))

"""
tt = listToTree([37,-34,-48,None,-100,-100,48,None,None,None,None,-54,None,-71,-22,None,None,None,8])

                    37
           -34                    -48
    None        -100        -100        48
None   None  None   None  -54   None  -71   -22


"""