1def distanceK(root, target, k):
2 parent_map = {}
3
4 def build_parent(node, parent):
5 if not node:
6 return
7 if parent:
8 parent_map[node] = parent
9 build_parent(node.left, node)
10 build_parent(node.right, node)
11
12 build_parent(root, None)
13
14 # BFS from target
15 queue = [target]
16 visited = {target}
17 dist = 0
18
19 while queue and dist < k:
20 level_size = len(queue)
21
22 for i in range(level_size):
23 current = queue.pop(0)
24
25 # Get neighbors: left, right, parent
26 neighbors = []
27 if current.left:
28 neighbors.append(current.left)
29 if current.right:
30 neighbors.append(current.right)
31 if current in parent_map:
32 neighbors.append(parent_map[current])
33
34 for neighbor in neighbors:
35 if neighbor not in visited:
36 visited.add(neighbor)
37 queue.append(neighbor)
38
39 dist += 1
40
41 result = [node.val for node in queue]
42 return result