forked from striver79/FreeKaTreeSeries
-
Notifications
You must be signed in to change notification settings - Fork 0
/
L30_nodesAtADistanceKJava
52 lines (52 loc) · 2.15 KB
/
L30_nodesAtADistanceKJava
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
class Solution {
private void markParents(TreeNode root, Map<TreeNode, TreeNode> parent_track, TreeNode target) {
Queue<TreeNode> queue = new LinkedList<TreeNode>();
queue.offer(root);
while(!queue.isEmpty()) {
TreeNode current = queue.poll();
if(current.left != null) {
parent_track.put(current.left, current);
queue.offer(current.left);
}
if(current.right != null) {
parent_track.put(current.right, current);
queue.offer(current.right);
}
}
}
public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
Map<TreeNode, TreeNode> parent_track = new HashMap<>();
markParents(root, parent_track, root);
Map<TreeNode, Boolean> visited = new HashMap<>();
Queue<TreeNode> queue = new LinkedList<TreeNode>();
queue.offer(target);
visited.put(target, true);
int curr_level = 0;
while(!queue.isEmpty()) { /*Second BFS to go upto K level from target node and using our hashtable info*/
int size = queue.size();
if(curr_level == k) break;
curr_level++;
for(int i=0; i<size; i++) {
TreeNode current = queue.poll();
if(current.left != null && visited.get(current.left) == null) {
queue.offer(current.left);
visited.put(current.left, true);
}
if(current.right != null && visited.get(current.right) == null ) {
queue.offer(current.right);
visited.put(current.right, true);
}
if(parent_track.get(current) != null && visited.get(parent_track.get(current)) == null) {
queue.offer(parent_track.get(current));
visited.put(parent_track.get(current), true);
}
}
}
List<Integer> result = new ArrayList<>();
while(!queue.isEmpty()) {
TreeNode current = queue.poll();
result.add(current.val);
}
return result;
}
}