leetcode-grind/problems/993/main.py

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

class Solution:
    
    def __init__(self):
        self.xPar = None
        self.xDepth = None
        self.yPar = None
        self.yDepth = None
        
    def isCousins(self, root: TreeNode, x: int, y: int) -> bool:
        # descend the binary tree and find the two values, set their parents
        # exit and then compare
        # because ints are unique and the nodes do exist there
        # shouldn't be too many edge cases I don't think
        self.descendTreeFinding(root, x, y, None, 0)
        return self.xDepth == self.yDepth and self.xPar != self.yPar
    
    def descendTreeFinding(self, tree, x, y, parent, depth):
        if tree is None:
            return
        
        if tree.val == x:
            self.xPar = parent
            self.xDepth = depth
            
        if tree.val == y:
            self.yPar = parent
            self.yDepth = depth
            
        if self.yPar is not None and self.xPar is not None:
            return
        else:
            self.descendTreeFinding(tree.left, x, y, tree.val, depth + 1)
            self.descendTreeFinding(tree.right, x, y, tree.val, depth + 1)

s = Solution()
print("Expected: true")
tree = TreeNode(1)
left = TreeNode(2)
left.right = TreeNode(4)
tree.left = left
right = TreeNode(3)
right.right = TreeNode(5)
tree.right = right

print("Got:", s.isCousins(tree, 4, 5))