leetcode-solution/Python/0993-cousins-in-binary-tree.py at main · hogan-tech/leetcode-solution · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
from collections import defaultdict
from typing import Optional


class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


class Solution:
    def isCousins(self, root: Optional[TreeNode], x: int, y: int) -> bool:
        cousinDict = defaultdict(tuple)

        if not root:
            return

        def traverse(node: Optional[TreeNode], level: int, parentVal: int):
            if node.val:
                cousinDict[node.val] = (level, parentVal)
            if node.left:
                traverse(node.left, level + 1, node.val)
            if node.right:
                traverse(node.right, level + 1, node.val)

        traverse(root, 0, float('inf'))
        return cousinDict[x][0] == cousinDict[y][0] and cousinDict[x][1] != cousinDict[y][1]


root = TreeNode(1)
root.left = TreeNode(2)
root.right = TreeNode(3)
root.left.left = TreeNode(4)
x = 4
y = 3
print(Solution().isCousins(root, x, y))