#733493 (C++17) No.1796 木上のクーロン

• 問題: No.1796 木上のクーロン
• ユーザー: Nachia
• 提出日時: 2022-01-22 18:18:31
• 言語: C++17 (gcc 13.3.0 + boost 1.87.0)
• 結果: AC
• 実行時間: 1,990 ms / 10,000 ms
• コード長: 8,453 bytes
• コンパイル時間: 1,501 ms
• コンパイル使用メモリ: 101,904 KB
• 最終ジャッジ日時: 2025-01-27 14:52:20
• ジャッジサーバーID （参考情報）: judge4 / judge4

ファイルパターン	結果
sample	AC * 3
other	AC * 34

ソースコード

#include <vector>
#include <utility>

namespace nachia{

struct AdjacencyList{
public:
    struct AdjacencyListRange{
        using iterator = typename std::vector<int>::const_iterator;
        iterator begi, endi;
        iterator begin() const { return begi; }
        iterator end() const { return endi; }
        int size() const { return (int)std::distance(begi, endi); }
    };
private:
    int mn;
    std::vector<int> E;
    std::vector<int> I;
public:
    AdjacencyList(int n, std::vector<std::pair<int,int>> edges, bool rev){
        mn = n;
        std::vector<int> buf(n+1, 0);
        for(auto [u,v] : edges){ ++buf[u]; if(rev) ++buf[v]; }
        for(int i=1; i<=n; i++) buf[i] += buf[i-1];
        E.resize(buf[n]);
        for(int i=(int)edges.size()-1; i>=0; i--){
            auto [u,v] = edges[i];
            E[--buf[u]] = v;
            if(rev) E[--buf[v]] = u;
        }
        I = std::move(buf);
    }
    AdjacencyList(const std::vector<std::vector<int>>& edges = {}){
        int n = mn = edges.size();
        std::vector<int> buf(n+1, 0);
        for(int i=0; i<n; i++) buf[i+1] = buf[i] + edges[i].size();
        E.resize(buf[n]);
        for(int i=0; i<n; i++) for(int j=0; j<(int)edges[i].size(); j++) E[buf[i]+j] = edges[i][j];
        I = std::move(buf);
    }
    AdjacencyListRange operator[](int u) const {
        return AdjacencyListRange{ E.begin() + I[u], E.begin() + I[u+1] };
    }
    int num_vertices() const { return mn; }
};

} // namespace nachia


namespace nachia{

struct CentroidDecomposition {
    int n;
    AdjacencyList E;
    std::vector<int> cdep;
    std::vector<int> cp;
    std::vector<int> cbfs;
    int maxdep;
    CentroidDecomposition(AdjacencyList edges) : E(std::move(edges)){
        n = E.num_vertices();
        std::vector<int> Z(n, 1);
        {
            std::vector<int> P(n, -1);
            std::vector<int> I = { 0 };
            for(int i=0; i<(int)I.size(); i++){
                int p = I[i];
                for (int e : E[p]) if (P[p] != e) {
                    P[e] = p;
                    I.push_back(e);
                }
            }
            for (int i=n-1; i>=1; i--) Z[P[I[i]]] += Z[I[i]];
        }
        cp.assign(n, -1);
        cdep.assign(n, 0);
        std::vector<std::pair<int, int>> I = { {0,-1} };
        for(int i=0; i<(int)I.size(); i++){
            int s = I[i].first;
            int par = I[i].second;
            while (true) {
                int nx = -1;
                for (int e : E[s]) if (Z[e] * 2 > Z[s]) nx = e;
                if (nx == -1) break;
                Z[s] -= Z[nx];
                Z[nx] += Z[s];
                s = nx;
            }
            cbfs.push_back(s);
            Z[s] = 0;
            if (par != -1) {
                cdep[s] = cdep[par] + 1;
                cp[s] = par;
            }
            for (int e : E[s]) if (Z[e] != 0) {
                I.push_back(std::make_pair(e, s));
            }
        }
        maxdep = 0;
        for (int a : cdep) maxdep = std::max(maxdep, a);
    }

    struct BFSUnit {
        std::vector<int> I;
        std::vector<int> P;
    };
    BFSUnit bfs_layer(int s, int layer) {
        BFSUnit res;
        if (cdep[s] < layer) return res;
        res.I.push_back(s);
        res.P.push_back(-1);
        for(int i=0; i<(int)res.I.size(); i++){
            int p = res.I[i];
            for (int e : E[p]) if (res.P[i] != e) {
                if (cdep[e] < layer) continue;
                res.I.push_back(e);
                res.P.push_back(p);
            }
        }
        return res;
    }
};

} // namespace nachia



#include <iostream>
#include <algorithm>
using namespace std;
using i64 = long long;
using u64 = unsigned long long;
using i32 = int;
using u32 = unsigned int;
#define rep(i,n) for(int i=0; i<int(n); i++)



// a^i mod M
template<u32 MOD>
u32 powm(u32 a,u64 i) {
    if(i == 0) return 1;
    u32 r = powm<MOD>((u64)a*a%MOD,i/2);
    if(i&1) r = (u64)r*a%MOD;
    return r;
}
 

template<u32 MOD, u32 g>
void NTT(vector<u32>& A){
    int N = 1;
    while (N < (int)A.size()) N *= 2;
    static vector<u32> exp_i_revbit_diff = { (u32)powm<MOD>(g, (MOD - 1) >> 2) };
    for(int i=exp_i_revbit_diff.size(); (1<<(i+1))<N; i++){
        u64 diffdiff = powm<MOD>(g, (MOD - 1) / 2 + ((MOD - 1) >> (i+2)) * 3);
        exp_i_revbit_diff.push_back(diffdiff);
    }
    for (int i = 1; i < N; i <<= 1) {
        u64 q = 1;
        int maxk = N / i / 2;
        for (int j = 0; j < i; j++) {
            int off = j * maxk * 2;
            for (int k = off; k < off + maxk; k++) {
                u32 l = A[k];
                u32 r = A[k + maxk] * q % MOD;
                A[k] = min(l + r - MOD, l + r);
                A[k + maxk] = min(l - r, l + MOD - r);
            }
            if(j+1 != i) for(int d=0; ; d++) if(!(j&(1<<d))){
                q = q * exp_i_revbit_diff[d] % MOD;
                break;
            }
        }
    }
    for (int i = 0, j = 0; j < N; j++) {
        if (i < j) swap(A[i], A[j]);
        for (int k = N >> 1; k > (i ^= k); k >>= 1);
    }
}



const u64 MOD = 998244353;
const u64 NTTzeta = 3;


int N;
vector<u64> Q;
nachia::AdjacencyList E;

u64 k0;
vector<u64> inv_mod;
vector<u64> C;
vector<vector<u32>> nttC;
vector<u64> inv_ntt_size;

vector<u64> ans;


vector<int> bfsbuf_dist;
void get_bfsbuf_dist(nachia::CentroidDecomposition::BFSUnit tree){
    bfsbuf_dist[tree.I.front()] = 0;
    for(int i=1; i<(int)tree.I.size(); i++){
        bfsbuf_dist[tree.I[i]] = bfsbuf_dist[tree.P[i]] + 1;
    }
}

vector<u32> sigma(const vector<u32>& B){
    int n = B.size() + 2;
    int Z = 1, d = 0;
    while(Z < n){ Z *= 2; d++; }
    vector<u32> revB(Z*2, 0);
    rep(i, B.size()) revB[Z-i] = B[i];
    NTT<MOD, NTTzeta>(revB);
    rep(i, Z*2) revB[i] = (u64)revB[i] * nttC[d+1][i] % MOD;
    NTT<MOD, NTTzeta>(revB);
    reverse(revB.begin() + 1, revB.end());
    rep(i, n) revB[i] = (u64)revB[i+Z] * inv_ntt_size[d+1] % MOD;
    revB.resize(n);
    return revB;
}

vector<u32> sigma_tree(const nachia::CentroidDecomposition::BFSUnit& tree){
    vector<u32> dist_freq(tree.I.size(), 0);
    get_bfsbuf_dist(tree);
    for(int p : tree.I){
        dist_freq[bfsbuf_dist[p]] += Q[p];
        if(dist_freq[bfsbuf_dist[p]] >= MOD) dist_freq[bfsbuf_dist[p]] -= MOD;
    }
    return sigma(dist_freq);
}


int main() {
    cin >> N;
    Q.resize(N);
    rep(i,N) cin >> Q[i];
    {
        vector<pair<int,int>> edges(N-1);
        rep(i,N-1){
            int u,v; cin >> u >> v; u--; v--;
            edges[i] = make_pair(u,v);
        }
        E = nachia::AdjacencyList(N, edges, true);
    }
    
    int max_ntt_size_log = 0;
    while((1 << max_ntt_size_log) < N + 6) max_ntt_size_log++;
    max_ntt_size_log++;
    int max_ntt_size = 1 << max_ntt_size_log;

    k0 = 1;
    for(int i=1; i<=N; i++) k0 = k0 * i % MOD;
    k0 = k0 * k0 % MOD;
    inv_mod.assign(max_ntt_size+1, 1);
    for(int i=2; i<=max_ntt_size; i++) inv_mod[i] = MOD - MOD / i * inv_mod[MOD % i] % MOD;
    C.resize(max_ntt_size);
    for(int i=0; i<max_ntt_size; i++) C[i] = k0 * inv_mod[i+1] % MOD * inv_mod[i+1] % MOD;

    nttC.resize(max_ntt_size_log+1);
    inv_ntt_size.resize(max_ntt_size_log+1);
    for(int d=0; d<=max_ntt_size_log; d++){
        nttC[d].resize(1<<d);
        for(int i=0; i<1<<d; i++) nttC[d][i] = C[i] % MOD;
        NTT<MOD, NTTzeta>(nttC[d]);
        inv_ntt_size[d] = powm<MOD>(1<<d, MOD-2);
    }

    bfsbuf_dist.assign(N, 0);
    ans.assign(N, 0);

    auto centroid_decomposition = nachia::CentroidDecomposition(E);
    for(int s=0; s<N; s++){
        int dep_s = centroid_decomposition.cdep[s];

        auto bfs_s = centroid_decomposition.bfs_layer(s, dep_s);
        vector<u32> sigma_s = sigma_tree(bfs_s);
        
        for(int nx : E[s]) if(centroid_decomposition.cdep[nx] > dep_s){
            nachia::CentroidDecomposition::BFSUnit bfs_nx = centroid_decomposition.bfs_layer(nx, dep_s+1);
            vector<u32> sigma_nx = sigma_tree(bfs_nx);
            for(int p : bfs_nx.I){
                int d = bfsbuf_dist[p] + 1;
                ans[p] = (ans[p] + (sigma_s[d] + MOD - sigma_nx[d+1])) % MOD;
            }
        }

        ans[s] = (ans[s] + sigma_s[0]) % MOD;
    }

    rep(p,N) cout << ans[p] << "\n";
    return 0;
}

struct ios_do_not_sync{
    ios_do_not_sync(){
        std::ios::sync_with_stdio(false);
        std::cin.tie(nullptr);
    }
} ios_do_not_sync_instance;