Problem

Given the root of a binary tree, the value of a target node target, and an integer k, return an array of the values of all nodes that have a distance k from the target node.

You can return the answer in any order.

Examples

graph TB;
    3 --- 5 & 1;
    5 --- 6 & 2;
    1 --- 0 & 8;
    2 --- 7 & 4;

style 5 fill:#FF9933
style 1 fill:#138808
style 7 fill:#138808
style 4 fill:#138808
  
1
2
3
4
5

Input:
root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, k = 2
Output:
 [7,4,1]
Explanation: The nodes that are a distance 2 from the target node (with value 5) have values 7, 4, and 1.

Example 2:

1
2
3
4

Input:
root = [1], target = 1, k = 3
Output:
 []

Similar Problem

Nodes at Distance k from Leaf Nodes in Binary Tree Nodes at Distance k from root in Binary Tree

Solution

Method 1 - BFS on Undirected Graph

  1. First create undirected graph using adj list
  2. Then start the BFS on target node. Whenever we find nodes at k distance, just add them to answer

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66

public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
	Map <TreeNode, List<TreeNode>> g = new HashMap();
	List<Integer> ans = new ArrayList<Integer>();

	if (root == null || k <0) {
		return ans;
	}

	buildGraph(g, root, null);

	if (!g.containsKey(target)) {
		return ans;
	}

	Set<TreeNode> visited = new HashSet<TreeNode>();
	Queue<TreeNode> q = new LinkedList<TreeNode>();
	q.add(target);
	visited.add(target);

	while (!q.isEmpty()) {
		int size = q.size();

		if (k == 0) {
			for (int i = 0; i <size; i++) {
				ans.add(q.poll().val);
			}

			return ans;
		}

		for (int i = 0; i <size; i++) {
			TreeNode node = q.poll();

			for (TreeNode nei: g.get(node)) {
				if (visited.contains(nei)){
					 continue;
				}
				visited.add(nei);
				q.add(nei);
			}
		}

		k--;
	}

	return ans;
}


private void buildGraph(Map <TreeNode, List<TreeNode>> g, TreeNode node, TreeNode parent) {
	if (node == null) {
		return;
	}

	if (!g.containsKey(node)) {
		g.put(node, new ArrayList<TreeNode>());

		if (parent != null) {
			g.get(node).add(parent);
			g.get(parent).add(node);
		}

		buildGraph(g, node.left, node);
		buildGraph(g, node.right, node);
	}
}

Complexity

  • Time: O(N)
  • Space: O(N)

Method 2 - BFS and Parents Map

Similar to above solution:

  1. First create parent map of all the nodes. We dont need to create undirected graph, as parent-child relation is easily handled by tree structure.
  2. Then start the BFS on target node. Keep on adding the child and parents to the queue during BFS till we have distance K.

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50

public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
	List<Integer> ans = new ArrayList<>();
	Map<TreeNode,TreeNode> parentsMap = new HashMap<>();
	buildParentMap(root,root,parentsMap);//first root is node and second root is its parent considered
	Set<TreeNode> visited = new HashSet<>(); //keep track of visited nodes
	Queue<TreeNode> q = new LinkedList<>();//for BFS
	q.add(target);  //instead of root here we are adding target as we need to find from that node
	visited.add(target); //also target node need to be visited first
	int level = 0; //as we are using level order traversal k == level then we find all the nodes in queue
	while(q.size()>0){
		int size = q.size();
		if(level == K) {
			return buildListFromQueue(q);//elements remaining in queue will be nodes at dist 
		}
		for(int i = 0; i < size; i++){                  //k from target
			TreeNode t = q.remove();
			if(t.left != null && !visited.contains(t.left)){
				q.add(t.left);
				visited.add(t.left);
			}
			if(t.right != null && !visited.contains(t.right)){
				q.add(t.right);
				visited.add(t.right);
			}
			TreeNode parent = parentsMap.get(t);
			if(!visited.contains(parent)){ //if parent is not visited 
				q.add(parent);//add parent to q
				visited.add(parent);   //mark it as visited
			}
		}
		level++; //we check levelwise and each level covers 1 unit distance 
	}
		return ans;
}
private List<Integer> buildListFromQueue(Queue<TreeNode> q){
	List<Integer> lst = new ArrayList<>();
	while(q.size()>0){
		lst.add(q.remove().val);
	}
	return lst;
}
private void buildParentMap(TreeNode root,TreeNode parent,Map<TreeNode,TreeNode> parentsMap){
	if(root == null) {
		return;
	}
		
	parentsMap.put(root, parent);
	buildParentMap(root.left,root, parentsMap); 
	buildParentMap(root.right,root, parentsMap);
}

Complexity

  • Time: O(N)
  • Space: O(N)