牛客網程式設計題 擴散II

牛客網原連結

題目描述

作為牛客王國的探險先鋒，牛牛來到了一片新的大陸。這是一個工業化程度很高的大陸，遍地都是工廠，有些工廠之間有管道相連。這片大陸一共有 \(n\) 個工廠，有 \(n\) 對工廠之間有管道相連，因為工廠之間需要合作，所以這 \(n-1\) 個管道保證任意兩個工廠都可以通過管道互相抵達。每根管道的長度都是1.（換言之，\(n\) 個工廠與 \(n-1\) 根管道組成了無根樹）。

牛牛發現，從這片大陸開始工業化以來，一共發生了 \(m\) 次汙染。
每次汙染用3個引數 \(x_i, y_i, z_i\)表示在第 \(x_i\) 個工廠發生了汙染，
影響了所有和第 \(x_i\)

一開始所有的工廠汙染指數為0.
現在牛牛想知道，在發生了這m次汙染之後，每個工廠的汙染指數是多少。

示例1

輸入

5,5,[1,2,3,3],[2,3,4,5],[1,2,3,4,5],[1,1,1,1,1],[1,2,3,4,5]

輸出

[3,6,14,7,8]

說明

在1，2，3，4，5分別發生了距離為1的汙染，他們的影響為分別為1，2，3，4，5
1發生的汙染影響了1，2
2發生的汙染影響了1，2，3
3發生的汙染影響了2，3，4，5
4發生的汙染影響了3，4
5發生的汙染影響了3，5
最終五個工廠汙染指數為[3,6,14,7,8]

備註

\(1<=n,m<=2e5, 1<=u,v,x_i<=n, 1<=y_i<=50, 1<=z_i<=1e4\)

第一個引數n代表工廠數量
第二個引數m代表汙染髮生次數
第三、四個引數vector u,v各自包含n-1個元素
第五、六、七個引數vector x,y,z各自包含m個元素，代表m次汙染

解題思路

最簡單的解法就是對每個 \(x_i\) 使用 bfs 增加汙染值，但是這種方法的時間複雜度最差為 \(O(n^2)\)，不能AC。
另一種方法：考慮以每個 \(x_i\) 擴散出來的路徑，在使用 \(O(n)\) 空間情況下，我們能否用上次擴散的路徑遞推得到下次路徑？並且在多個 \(x_i\)

下面舉個例子來說明這個過程：
假設這樣一個工廠連線圖，輸入的指令是 \(x=3, y=3, z=7\)。

用兩個陣列（A，B）來儲存路徑距離源點 \(3\) 的距離為奇數和偶數的點上的汙染值，這樣最終的結果就是兩個陣列相加即可。這兩個陣列並不固定哪個是奇數哪個是偶數，只是固定B一點包含傳播最遠的點。

1. 開始的時候，源點為 3, 所以 B[3] = 7;
2. 傳播一步得到新增加的點，B[2,4,5] = 7, A 就是原先的 B。
3. 。。。
   
   更新公式：

     C = A
     for i in range(1,n):
          for neig in grpha[i]:
                 C[i] += B[neig] - A[i]
     A = B
     B = C

不同的源點\(x\)之間可以共用一個數組，並不影響最終結果，不同的距離 \(y\) 意味這跌打更新的輪次不同，最長的 \(y\) 可能最先開始，每次迭代減少一個距離。所以按 \(y\) 的倒序處理每個源點。
綜上程式碼如下：
!!! 測試資料裡面有坑，好像有 \(x_i=0\)

#include <vector>
#include <algorithm>
#include <iostream>
using namespace std;
class Solution {
public:
    struct Pul{
        int x,y,z;
    };
    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>()); // gaphe 圖
        for(int i=0; i<u.size(); i++){
            g[u[i]].push_back(v[i]);
            g[v[i]].push_back(u[i]);
        }
        vector<Pul> ps;
        for(int i=0; i<x.size(); ++i){
            ps.push_back({x[i], y[i], z[i]});
        }
        sort(ps.begin(), ps.end(), [](Pul& a, Pul& b)->bool{ return a.y<b.y; } );
        vector<int> a(n+1, 0), b(n+1, 0);
        for(int sz=ps.back().y; sz>=0; --sz){ // 倒序處理
            vector<int> c(a);
            for(int i=0; i<=n; i++){
                for(auto& neig:g[i]) c[i] += b[neig] - a[i];
            }
            while(ps.size()){
                auto& tmp = ps.back();
                if(tmp.y < sz) break;
                c[tmp.x] += tmp.z;
                ps.pop_back();
            }
            swap(a,b);
            swap(b,c);
        }
        for(int i=0; i<=n; i++) a[i] += b[i];
        a.erase(a.begin());
        return a;
    }
};