Cousins in Binary Tree Python Solution

lcpy gen -n 993   # by problem number
lcpy gen -s cousins_in_binary_tree   # by problem name

Input: root = [1,2,3,4], x = 4, y = 3
Output: false

Input: root = [1,2,3,null,4,null,5], x = 5, y = 4
Output: true

Input: root = [1,2,3,null,4], x = 2, y = 3
Output: false

from leetcode_py import TreeNode


class Solution:
    # Time: O(n) - single traversal of the tree
    # Space: O(h) - recursion stack, h is the tree height
    def is_cousins(self, root: TreeNode[int] | None, x: int, y: int) -> bool:
        info: dict[int, tuple[int, int | None]] = {}

        def dfs(node: TreeNode[int] | None, parent: int | None, depth: int) -> None:
            if node is None:
                return
            info[node.val] = (depth, parent)
            dfs(node.left, node.val, depth + 1)
            dfs(node.right, node.val, depth + 1)

        if root is None:
            return False
        dfs(root, None, 0)
        dx, px = info[x]
        dy, py = info[y]
        return dx == dy and px != py