All Nodes Distance K in Binary Tree LeetCode solution in Java

1
class Solution {
2
  public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
3
    HashMap<TreeNode, TreeNode> map = new HashMap<>();
4
    get_parent(root, map);
5
    Queue<TreeNode> q = new LinkedList<>();
6
    q.add(target);
7
    int distance = 0;
8
    HashSet<TreeNode> visited = new HashSet<>();
9
    visited.add(target);
10
    while (!q.isEmpty()) {
11
      if (distance == k) break;
12
      distance++;
13
      int size = q.size();
14
      for (int i = 0; i < size; i++) {
15
        TreeNode current = q.poll();
16
        if (current.left != null && !visited.contains(current.left)) {
17
          q.add(current.left);
18
          visited.add(current.left);
19
        }
20
        if (current.right != null && !visited.contains(current.right)) {
21
          q.add(current.right);
22
          visited.add(current.right);
23
        }
24
        if (map.containsKey(current) && !visited.contains(map.get(current))) {
25
          q.add(map.get(current));
26
          visited.add(map.get(current));
27
        }
28
      }
29
    }
30
    List<Integer> ans = new ArrayList<>();
31
    while (!q.isEmpty()) ans.add(q.poll().val);
32
    return ans;
33
  }
34

35
  public void get_parent(TreeNode root, HashMap<TreeNode, TreeNode> map) {
36
    Queue<TreeNode> q = new LinkedList<>();
37
    q.add(root);
38
    while (!q.isEmpty()) {
39
      int size = q.size();
40
      for (int i = 0; i < size; i++) {
41
        TreeNode current = q.poll();
42
        if (current.left != null) {
43
          map.put(current.left, current);
44
          q.add(current.left);
45
        }
46
        if (current.right != null) {
47
          map.put(current.right, current);
48
          q.add(current.right);
49
        }
50
      }
51
    }
52
  }
53
}