Cousins in Binary Tree LeetCode solution in Python

1
# Definition for a binary tree node.
2
# class TreeNode:
3
#     def __init__(self, val=0, left=None, right=None):
4
#         self.val = val
5
#         self.left = left
6
#         self.right = right
7
class Solution:
8
    ans = False
9

10
    def isCousins(self, root: Optional[TreeNode], x: int, y: int) -> bool:
11
        def dfs(node, depth):
12
            if self.ans or not node:
13
                return 0
14
            if node.val == x or node.val == y:
15
                return depth
16
            l = dfs(node.left, depth + 1)
17
            r = dfs(node.right, depth + 1)
18
            if not (l and r):
19
                return l or r
20
            if l == r and l != depth + 1:
21
                self.ans = True
22
            return 0
23

24
        dfs(root, 0)
25
        return self.ans