863. All Nodes Distance K in Binary Tree | 1663

 1# Definition for a binary tree node.
 2# class TreeNode:
 3#     def __init__(self, x):
 4#         self.val = x
 5#         self.left = None
 6#         self.right = None
 7
 8
 9class Solution:
10    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:
11        # intuit: find a way to store self.parent
12        # since Node.val are unique, a simple dict should work
13        # we also can't look back: e.g., 5 -> 2 -> 5
14        # therefore we can also use a visited set() to prevent looping
15        # bfs is easier for distance problems
16
17        parent = {}
18        visited = set()
19
20        # first traversal map parent{}
21        q = deque([(root, None)])  # node, node.parent
22
23        while q:
24            node, node_parent = q.popleft()
25            parent[node.val] = node_parent
26            if node.left:
27                q.append((node.left, node))
28            if node.right:
29                q.append((node.right, node))
30
31        q = deque([(target, 0)])  # node, distance
32        result = []
33
34        while q:
35            node, dist = q.popleft()
36            if dist == k:
37                result.append(node.val)
38                continue
39            visited.add(node.val)
40
41            if node.left and node.left.val not in visited:
42                q.append((node.left, dist + 1))
43            if node.right and node.right.val not in visited:
44                q.append((node.right, dist + 1))
45            parent_node = parent[node.val]
46            if parent_node and parent_node.val not in visited:
47                q.append((parent_node, dist + 1))
48
49        return result