题解 | #扩散#_牛客博客

思路：

题目的主要信息：

$n$ 个节点， $n-1$ 边使之任意两两都有相连，这是一棵树
每个节点初始值为0，一共m次增加数字的机会，会使给出的节点及与之直接相连的节点值加1
求最后的每个节点的数字

方法一：模拟过程
具体做法：
用构建图的方式，构建这棵树。遍历m次增加的机会，每次找到要增加节点的下标，将其值加1，然后遍历与下标相连的其他节点，值都加1.
这里的res本来应该是哈希表的，但是因为数组的下标可以与节点名称一一对应，因此直接用数组。
图片说明

class Solution {
public:
    vector<int> solve(int n, int m, vector<int>& u, vector<int>& v, vector<int>& q) {
        vector<int> res(n, 0);
        vector<vector<int> > G(n + 1);
        for(int i = 0; i < u.size(); i++){ //构建图
            G[u[i]].push_back(v[i]);
            G[v[i]].push_back(u[i]);
        }
        for(int i = 0; i < m; i++){
            int temp = q[i];
            res[temp - 1]++; //当前工厂增加
            for(int j = 0; j < G[temp].size(); j++) //与之相连的工厂增加
                res[G[temp][j] - 1]++;
        }
        return res;
    }
};

复杂度分析：

时间复杂度： $O(mn)$ ，最坏情况一个根节点，全是子节点，所有增加机会都是根节点，就是n个节点乘上m次机会
空间复杂度： $O(n^2)$ ，辅助数组G的大小，最坏情况为邻接矩阵，其中res是返回函数必要空间

方法二：各自统计
具体做法：
一共是两种增加值的情况，一种是被选中了，自己增加；另一种则是因为我旁边有直接相连的，被带动增加了。我们可以用两个数组分别记录每个节点这两种情况的各自增加了多少，当然必须先算第一个数组。然后两个数组相同下标值相加即是答案。
图片说明

class Solution {
public:
    vector<int> solve(int n, int m, vector<int>& u, vector<int>& v, vector<int>& q) {
        vector<int> res(n, 0);
        vector<int> add1(n + 1, 0); //记录自己增加的数量
        vector<int> add2(n + 2, 0); //记录因与别的节点直接相连增加的数
        for(int i = 0; i < m; i++)  //自己
            add1[q[i]] += 1;
        for(int i = 0; i < n - 1; i++){ //直接相连
            add2[u[i]] += add1[v[i]];
            add2[v[i]] += add1[u[i]];
        }
        for(int i = 1; i <=n; i++)  //二者相加即是答案
            res[i - 1] = add1[i] + add2[i];
        return res;
    }
};

复杂度分析：

时间复杂度： $O(m+n)$ ，循环是分开的，第一次为 $O(m)$ ，后面两次都是 $O(n)$
空间复杂度： $O(n)$ ，两个辅助数组add1和add2，res为返回函数的必要空间