All Nodes Distance K in Binary Tree LeetCode solution in C++

1
void makeParent(TreeNode *root, unordered_map<TreeNode *, TreeNode *> &parent) {
2
  queue<TreeNode *> q;
3
  q.push(root);
4

5
  while (q.empty() == false) {
6
    TreeNode *curr = q.front();
7
    q.pop();
8

9
    if (curr->left != NULL) {
10
      parent[curr->left] = curr;
11
      q.push(curr->left);
12
    }
13

14
    if (curr->right != NULL) {
15
      parent[curr->right] = curr;
16
      q.push(curr->right);
17
    }
18
  }
19
}
20

21
class Solution {
22
 public:
23
  void solve(TreeNode *target, unordered_map<TreeNode *, TreeNode *> &parent,
24
             unordered_map<TreeNode *, bool> &visited, vector<int> &res, int k) {
25
    if (k == 0) {
26
      res.push_back(target->val);
27
      return;
28
    }
29

30
    queue<TreeNode *> q;
31
    int distance = 0;
32

33
    q.push(target);
34
    visited[target] = 1;
35

36
    while (q.empty() == false) {
37
      int count = q.size();
38
      distance++; // increment distance value for current level of our level
39
                  // order traversal
40

41
      for (int i = 0; i < count; i++) {
42
        TreeNode *curr = q.front();
43
        q.pop();
44

45
        if (curr->left != NULL && visited[curr->left] == false) {
46
          visited[curr->left] = true;
47
          q.push(curr->left);
48

49
          if (distance == k) // when we reach the node at distance = k, push it to res
50
            res.push_back(curr->left->val);
51
        }
52

53
        if (curr->right != NULL && visited[curr->right] == false) {
54
          visited[curr->right] = true;
55
          q.push(curr->right);
56

57
          if (distance == k) // when we reach the node at distance = k, push it to res
58
            res.push_back(curr->right->val);
59
        }
60

61
        if (parent[curr] && visited[parent[curr]] == false) {
62
          visited[parent[curr]] = true;
63
          q.push(parent[curr]);
64

65
          if (distance == k) // when we reach the node at distance = k, push it to res
66
            res.push_back(parent[curr]->val);
67
        }
68
      }
69

70
      if (distance == k) // if in curr iteration all k distances nodes have been
71
                         // pushed, we break
72
        break;
73
    }
74
  }
75

76
  vector<int> distanceK(TreeNode *root, TreeNode *target, int k) {
77
    unordered_map<TreeNode *, TreeNode *> parent;
78
    makeParent(root, parent);
79

80
    unordered_map<TreeNode *, bool> visited;
81
    vector<int> ans;
82

83
    solve(target, parent, visited, ans, k); // begin from target Node
84
    return ans;
85
  }