863. All Nodes Distance K in Binary Tree

Posted on Edited on In

863. All Nodes Distance K in Binary Tree

Description

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.

Example 1:

1
2
3
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
Input: root = [1], target = 1, k = 3
Output: []

Constraints:

  • The number of nodes in the tree is in the range [1, 500].
  • 0 <= Node.val <= 500
  • All the values Node.val are unique .
  • target is the value of one of the nodes in the tree.
  • 0 <= k <= 1000

Hints/Notes

  • binary tree

Solution

Language: C++

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
/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    map<TreeNode*, TreeNode*> parent;

    vector<int> distanceK(TreeNode* root, TreeNode* target, int k) {
        traverse(root);
        int distance = 0;
        queue<TreeNode*> q;
        q.push(target);
        unordered_set<int> s;
        vector<int> res;
        while (!q.empty()) {
            int n = q.size();
            for (int i = 0; i < n; i++) {
                TreeNode* cur = q.front();
                q.pop();
                if (distance == k) {
                    res.push_back(cur->val);
                }
                s.insert(cur->val);
                if (cur->left && !s.contains(cur->left->val)) {
                    q.push(cur->left);
                }
                if (cur->right && !s.contains(cur->right->val)) {
                    q.push(cur->right);
                }
                if (parent[cur] && !s.contains(parent[cur]->val)) {
                    q.push(parent[cur]);
                }
            }
            distance++;
            if (distance > k) {
                break;
            }
        }
        return res;
    }

    void traverse(TreeNode* root) {
        if (!root) {
            return;
        }
        if (root->left) {
            parent[root->left] = root;
            traverse(root->left);
        }
        if (root->right) {
            parent[root->right] = root;
            traverse(root->right);
        }
    }
};