CODEWITHSUNDEEP
pythontreebinary-tree
younger
younger

Reputation: 35

Python find lowest common ancestor of two nodes in a binary tree if not all of these nodes in the tree

I understand how to solve the question if these two nodes must in the Binary Tree, but what if they do not have to be in the tree? If only one or none of these nodes in the tree, return None.

Here is my code:

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

class Solution(object):
    def lowestCommonAncestor(self, root, p, q):
        """
        :type root: TreeNode
        :type p: TreeNode
        :type q: TreeNode
        :rtype: TreeNode
        """
        [root,count] = self.findAncestor(root,p,q,0)
        if count == 2:
            return root
        else:
            return None
    def findAncestor(self,root,p,q,count):
        if not root:
            return None, count
        left,left_count = self.findAncestor(root.left, p, q,count)
        right,right_count = self.findAncestor(root.right,p,q,count)
        if root == p or root == q:
            return root,count+1
        if left and right:
            return root,count
        elif left:
            return left,left_count
        else:
            return right,right_count

but I keep getting incorrect answer. Anyone know how to fix it based on my code? thanks!

Upvotes: 1

Views: 2603

Answers (2)

Wangdu Lin
Wangdu Lin

Reputation: 913

Base on kitt's solution, I test his solution on lintCode problem 578, but it did not passed. The issue happens at counting condition, which should check one more time with those input two node. So I redesign a new solution which passed lintcode test, also with better reading logic.

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


class Solution:
    """
    @param: root: The root of the binary tree.
    @param: A: A TreeNode
    @param: B: A TreeNode
    @return: Return the LCA of the two nodes.
    """
    count = 0

    def lowestCommonAncestor3(self, root, A, B):
        result = self.lca(root, A, B)
        return result if self.count == 2 else None

    def lca(self, root, A, B):
        if not root:
            return None

        for node in [A, B]:
            if root == node:
                self.count += 1

        left = self.lca(root.left, A, B)
        right = self.lca(root.right, A, B)

        if root in (A, B) or left and right:
            return root

        if left:
            return left

        if right:
            return right

        return None

Upvotes: 0

kitt
kitt

Reputation: 120

We can count target node number and if it's 2, then we know both nodes are in the tree.

class Solution(object):
    def lowestCommonAncestor(self, root, p, q):
        """
        :type root: TreeNode
        :type p: TreeNode
        :type q: TreeNode
        :rtype: TreeNode
        """
        self.count = 0
        node = self.find(root, p, q)
        return node if self.count == 2 else None

    def find(self, node, p, q):
        if not node:
            return None
        if node in (p, q):
            self.count += 1
        left = self.find(node.left, p, q)
        right = self.find(node.right, p, q)
        return node if node in (p, q) or left and right else left or right

Upvotes: 3

Related Questions