Lintcode 88. Lowest Common Ancestor of a Binary Tree (Medium) (Python)

Lowest Common Ancestor of a Binary Tree

Description:

Given the root and two nodes in a Binary Tree. Find the lowest common ancestor(LCA) of the two nodes.

The lowest common ancestor is the node with largest depth which is the ancestor of both nodes.

Example
For the following binary tree:

4
/ \
3 7
/ \
5 6
LCA(3, 5) = 4

LCA(5, 6) = 7

LCA(6, 7) = 7

Notice
Assume two nodes are exist in tree.

Code:

"""
Definition of TreeNode:
class TreeNode:
    def __init__(self, val):
        self.val = val
        self.left, self.right = None, None
"""


class Solution:
    """
    @param: root: The root of the binary search tree.
    @param: A: A TreeNode in a Binary.
    @param: B: A TreeNode in a Binary.
    @return: Return the least common ancestor(LCA) of the two nodes.
    """
    def lowestCommonAncestor(self, root, A, B):
        # write your code here
        if not root or root==A or root==B :
            return root
        l = self.lowestCommonAncestor(root.left, A, B)
        r = self.lowestCommonAncestor(root.right, A, B)
        if l and r:
            return root
        if l:
            return l
        elif r:
            return r