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

算法思想一：直接模拟

解题思路：

主要是通过构建图的方式，构建一棵树。

遍历m次增加的机会，每次找到要增加节点的下标，将其值加1，然后遍历与下标相连的其他节点，值都加1.

代码展示：

C++版本

class Solution {
public:
    /**
     * 代码中的类名、方法名、参数名已经指定，请勿修改，直接返回方法规定的值即可
     *
     * 
     * @param n int整型 
     * @param m int整型 
     * @param u int整型vector 
     * @param v int整型vector 
     * @param q int整型vector 
     * @return int整型vector
     */
    vector<int> solve(int n, int m, vector<int>& u, vector<int>& v, vector<int>& q) {
        // write code here
        vector<int> res(n, 0);
        vector<vector<int> > tree(n + 1);
        for(int i = 0; i < u.size(); i++){ //构建图
            tree[u[i]].push_back(v[i]);
            tree[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 < tree[temp].size(); j++) //与之相连的工厂增加
                res[tree[temp][j] - 1]++;
        }
        return res;
    }
};

复杂度分析

时间复杂度 $O(mn)$ ：最坏情况一个根节点，全是子节点，所有增加机会都是根节点，就是n个节点乘上m次机会

空间复杂度 $O(n^2)$ ：辅助数组tree的大小，最坏情况为邻接矩阵，其中res是返回函数必要空间

算法思想二：分离计算

解题思路：

将题目分为是两种增加值的情况，

1、一种是被选中了，自己增加

2、另一种则是因为旁边有直接相连的，被带动增加了

因此可以用两个数组one、two分别记录每个节点这两种情况的各自增加了多少，然后两个数组相同下标值相加即是答案。
注：必须先算第一个数组

图解：

代码展示：

Python3

#
# 代码中的类名、方法名、参数名已经指定，请勿修改，直接返回方法规定的值即可
#
# 
# @param n int整型 
# @param m int整型 
# @param u int整型一维数组 
# @param v int整型一维数组 
# @param q int整型一维数组 
# @return int整型一维数组
#
class Solution:
    def solve(self , n , m , u , v , q ):
        # write code here
        res = [0] * n
        # 记录自己增加的数量
        one = [0] * (n+1)
        # 记录因与别的节点直接相连增加的数
        two = [0] * (n+2)
        # 自己
        for i in range(m):
            one[q[i]] += 1
        # 直接相连
        for i in range(n-1):
            two[u[i]] += one[v[i]]
            two[v[i]] += one[u[i]]
        # 二者相加即是答案
        for i in range(1, n+1):
            res[i-1] = one[i] + two[i]
        return res

复杂度分析

时间复杂度 $O(m+n)$ ：两次循环，第一次为 $O(m)$ ，后面两次都是 $O(n)$

空间复杂度 $O(n)$ ：两个辅助数组one和two，res为返回函数的必要空间