Mind palace

some thoughts

0%

863. All Nodes Distance K in Binary Tree

Posted on Edited on In Disqus:
Symbols count in article: 2.2k Reading time ≈ 2 mins.

O(n) time O(n) space
树转图+bfs
跟742. Closest Leaf in a Binary Tree类似

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
/**
 * 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:
    vector<int> distanceK(TreeNode* root, TreeNode* target, int K) {
        dfs(root);
        list<int> q{{target->val}};
        while (K-- > 0 && !q.empty()) {
            for (int i = q.size(); i > 0; --i) {
                int u = q.front(); q.pop_front();
                for (int v : g[u]) {
                    g[v].erase(u);
                    q.push_back(v);
                }
            }
        }
        return vector<int>(begin(q), end(q));
    }

    void dfs(TreeNode *root) {
        if (!root) return;
        if (root->left) {
            g[root->val].insert(root->left->val);
            g[root->left->val].insert(root->val);
        }
        if (root->right) {
            g[root->val].insert(root->right->val);
            g[root->right->val].insert(root->val);
        }
        dfs(root->left);
        dfs(root->right);
    }

    unordered_map<int, unordered_set<int>> g; // 邻接表
};
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
/**
 * 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:
    vector<int> distanceK(TreeNode* root, TreeNode* target, int K) {
        dfs(root);
        unordered_set<int> s{target->val};
        list<int> q{{target->val}};
        while (K-- > 0 && !q.empty()) {
            for (int i = q.size(); i > 0; --i) {
                int x = q.front(); q.pop_front();
                for (int y : g[x]) {
                    if (s.count(y)) continue;
                    s.insert(y);
                    q.push_back(y);
                }
            }
        }
        return vector<int>(begin(q), end(q));
    }

    void dfs(TreeNode *root) {
        if (!root) return;
        if (root->left) {
            g[root->val].push_back(root->left->val);
            g[root->left->val].push_back(root->val);
        }
        if (root->right) {
            g[root->val].push_back(root->right->val);
            g[root->right->val].push_back(root->val);
        }
        dfs(root->left);
        dfs(root->right);
    }

    unordered_map<int, vector<int>> g; // 邻接表
};