Leetcode 863 - All Nodes Distance K in Binary Tree

題目

Problem#

給你一個二元樹，和目標節點(target）以及一個正整數 k ，要你回傳所有與 target 距離 k 的節點

測資限制#

• number of nodes: $0 \le n \le 500$
• $0 \le k \le 1000$$
• node value: $0 \le val \le 500$

想法#

題目給的是 binary tree ，只能往 leaf 走，並不能往回走，所以可以將題目先轉成 graph，接著在上面遍歷，找距離 $k$ 的節點即可

• 時間複雜度: $\mathcal{O}(n)$
• 空間複雜度: $\mathcal{O}(n)$

AC Code#

Copy

class Solution {
public:
    vector<vector<int>> G;
    map<TreeNode*, int> N;
    map<int,TreeNode*> NN;
    int id = 0;
    int getId(TreeNode *n)
    {
        if(!N.count(n))
        {
            NN[id] = n;
            N[n] = id++;
            G.push_back({});
        }
        return N[n];
    }
    void build(TreeNode *n)
    {
        if(!n) return;
        if(n->left)
        {
            int a = getId(n), b = getId(n->left);
            G[a].push_back(b);
            G[b].push_back(a);
            build(n->left);
        }
        if(n->right)
        {
            int a = getId(n), b = getId(n->right);
            G[a].push_back(b);
            G[b].push_back(a);
            build(n->right);
        }
    }
    vector<int> ans;
    void bfs(int s, int k)
    {
        using P=pair<int,int>;
        vector<bool> vis(N.size());
        queue<P> q;
        q.emplace(s, 0);
        while(!q.empty())
        {
            auto [cur, dep] = q.front();
            q.pop();
 
            if(dep >= k)
            {
                if(cur != s)
                    ans.push_back(cur);
                continue;
            }
 
            for(auto to : G[cur])
            {
                if(!vis[to])
                {
                    vis[to] = true;
                    q.emplace(to, dep+1);
                }
            }
        }
    }
    vector<int> distanceK(TreeNode* r, TreeNode* t, int k)
    {
        if(k == 0)
            return {t->val};
        build(r);
        int s = getId(t);
        bfs(s, k);
        for_each(ans.begin(), ans.end(), [&](auto &i) { i = NN[i]->val; });
        return ans;
    }
};

view raw leetcode/863.cpp delivered with ❤ by emgithub

賞析#

有人作法只記每個點的 parent 這樣還是可以遍歷整棵樹