Closest Leaf in a Binary Tree Python Solution

lcpy gen -n 742   # by problem number
lcpy gen -s closest_leaf_in_a_binary_tree   # by problem name

Input: root = [1,3,2], k = 1
Output: 2
Explanation: Either 2 or 3 is the nearest leaf node to the target of 1.

Input: root = [1], k = 1
Output: 1
Explanation: The nearest leaf node is the root node itself.

Input: root = [1,2,3,4,null,null,null,5,null,6], k = 2
Output: 3
Explanation: The leaf node with value 3 (and not the leaf node with value 6) is nearest to the node with value 2.

from collections import deque

from leetcode_py import TreeNode


class Solution:
    # Time: O(n)
    # Space: O(n)
    def find_closest_leaf(self, root: TreeNode[int] | None, k: int) -> int:
        assert root is not None
        parent: dict[int, TreeNode[int] | None] = {id(root): None}
        q: deque[TreeNode[int]] = deque([root])
        target: TreeNode[int] | None = None
        while q:
            node = q.popleft()
            if node.val == k:
                target = node
            for child in (node.left, node.right):
                if child is not None:
                    parent[id(child)] = node
                    q.append(child)
        assert target is not None
        dist: dict[int, int] = {id(target): 0}
        q = deque([target])
        while q:
            node = q.popleft()
            if node.left is None and node.right is None:
                return node.val
            for nxt in (parent[id(node)], node.left, node.right):
                if nxt is not None and id(nxt) not in dist:
                    dist[id(nxt)] = dist[id(node)] + 1
                    q.append(nxt)
        return root.val