First build parent pointers with DFS, then BFS from the target to distance k.
function distanceK(root, target, k):
parent = {}
function buildParent(node, par):
if node is null:
return
parent[node] = par
buildParent(node.left, node)
buildParent(node.right, node)
buildParent(root, null)
visited = set()
queue = [target]
visited.add(target)
dist = 0
while queue is not empty:
if dist == k:
return [node.val for node in queue]
nextLevel = []
for node in queue:
for neighbor in [node.left, node.right, parent[node]]:
if neighbor is not null and neighbor not in visited:
visited.add(neighbor)
nextLevel.append(neighbor)
queue = nextLevel
dist += 1
return []
time, space.