Given the root of a binary tree, a target node, and an integer k, return all the nodes that are distance k from the target node.
Input: root=[3,5,1,6,2,0,8,null,null,7,4], target=5, k=2 → Output: [7,4,1]Input: root=[1], target=1, k=3 → Output: []
Convert the tree to an undirected graph (tracking parent pointers). Then BFS from target node with distance k.
- DFS to build parent map: parent[node] = its parent.
- BFS from target node, tracking visited nodes.
- Expand each level; after k levels, collect all nodes in the queue.
import java.util.*;
class Solution {
Map<TreeNode, TreeNode> parent = new HashMap<>();
public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
buildParent(root, null);
Queue<TreeNode> q = new LinkedList<>();
Set<TreeNode> visited = new HashSet<>();
q.add(target); visited.add(target);
int dist = 0;
while (!q.isEmpty()) {
if (dist == k) {
List<Integer> res = new ArrayList<>();
for (TreeNode n : q) res.add(n.val);
return res;
}
int size = q.size();
for (int i = 0; i < size; i++) {
TreeNode n = q.poll();
for (TreeNode nb : new TreeNode[]{n.left, n.right, parent.get(n)}) {
if (nb != null && !visited.contains(nb)) { q.add(nb); visited.add(nb); }
}
}
dist++;
}
return new ArrayList<>();
}
private void buildParent(TreeNode node, TreeNode par) {
if (node == null) return;
parent.put(node, par);
buildParent(node.left, node); buildParent(node.right, node);
}
}
- Time Complexity: O(N)
- Space Complexity: O(N)