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

思路：

题目的主要信息：

n个节点，n-1条边使之连通，这就是一棵树（注意不一定是二叉树），每条边代表距离为1
一共m次污染，每次发生在数组元素x[i]，影响范围是与发生点距离不超过y[i]（发生点视为距离为0），影响范围所有节点污染指数增加z[i]
污染指数初始值全部为0，求m次污染发生后，每个节点的污染指数

方法一：排序+距离模拟
具体做法：
我们可以对所有污染的影响距离进行排序，然后以影响范围最远的开始，依次递减，直到为0，我们用三个数组相互交替记录之前增加的污染指数和本轮增加的污染指数，最后将二者相加即是结果。计算本轮的时候需要注意不能当前的距离不能小于本轮的距离，否则会重复，此时应该直接跳出循环，这里排序的意义在于可以在 $O(1)$ 复杂度之内pop掉尾数据。计算扩散的时候需要用到上一轮的数据，直接遍历无向图进行扩散，增加附近的即可，需要去重，即减掉已经上次增加的。

struct node{
    int x, y, z;
};
class Solution {
public:
    static bool comp(node& a, node& b){
        return a.y < b.y;
    }
    vector<int> solve(int n, int m, vector<int>& u, vector<int>& v, vector<int>& x, vector<int>& y, vector<int>& z) {
        vector<vector<int> > G(n + 1);
        vector<int> res;
        for(int i = 0; i < u.size(); i++){ //构建图
            G[u[i]].push_back(v[i]);
            G[v[i]].push_back(u[i]);
        }
        vector<node> polution;
        for(int i = 0; i < m; i++)
            polution.push_back({x[i], y[i], z[i]});
        sort(polution.begin(), polution.end(), comp); //按照污染影响距离排序
        vector<int> a(n + 1, 0), b(n + 1, 0);
        for(int i = polution.back().y; i >= 0; i--){
            vector<int> c = a; //记录本轮有多少污染指数
            for(int j = 0; j <= n; j++){  //之前的污染扩散
                for(int k = 0; k < G[j].size(); k++)
                    c[j] += b[G[j][k]] - a[j];
            }
            while(polution.size() != 0){
                auto temp = polution.back(); //本次污染
                if(temp.y < i) //最大距离小于当前距离，不用继续循环
                    break;
                c[temp.x] += temp.z; //增加
                polution.pop_back();
            }
            a = b; //依次往上叠加
            b = c; 
        }
        for(int i = 1; i <= n; i++)
           res.push_back(a[i] + b[i]);  //之前累加的加上本轮增加的
        return res;
    }
};

复杂度分析：

时间复杂度： $O(mlog_2m+max(x_i)*n)$ ，主体循环外有一个sort排序函数，外循环是从 $max(x_i)$ 到 $0$ ，内循环总共最多遍历 $n$ 次
空间复杂度： $O(n^2)$ ，最坏情况无向图构成邻接矩阵

方法二：dfs
具体做法：
我们可以用邻接矩阵构建一个无向图。然后遍历每次发生污染的初始节点，对其进行深度优先遍历，因为每条边代表距离为1，因此距离起始点的深度即距离。发生污染的点我们视作深度为0，利用dfs遍历到深度为y[i]为止，将其每个点的污染指数增加z[i]。
图片说明

class Solution {
public:
    void dfs(vector<vector<int>>& G, vector<int>& res, int prev, int x, int y, int z, int depth) {
        res[x] += z; //先加节点自身
        for(int i = 0; i < G[x].size(); i++) //遍历所有连接节点
            if (G[x][i] != prev)  //遍历非前序
                if (depth + 1 <= y)
                    dfs(G, res, x, G[x][i], y, z, depth + 1); //深度需要加1
    }
    vector<int> solve(int n, int m, vector<int>& u, vector<int>& v, vector<int>& x, vector<int>& y, vector<int>& z) {
        vector<vector<int> > G(n + 1);
        vector<int> res(n + 1, 0);
        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++) //对每一次污染节点进行深度优先遍历，距离即深度
            dfs(G, res, -1, x[i], y[i], z[i], 0); //对于每个节点前序节点为-1，初始深度为0
        res.erase(res.begin());
        return res;
    }
};

复杂度分析：

时间复杂度： $O(mn)$ ，一共发生 $m$ 次污染，最坏情况下每次污染的影响 $n$ 个节点
空间复杂度： $O(n^2)$ ，递归栈深度最坏情况下为每次污染的影响 $n$ 个节点且树为链表型，复杂度为 $O(n)$ ，构建的图最坏为邻接矩阵，复杂度为 $O(n^2)$ ，res属于返回函数必要空间