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

• 問題: No.1796 木上のクーロン
• ユーザー: Nachia
• 提出日時: 2022-10-07 23:10:38
• 言語: C++17 (gcc 13.3.0 + boost 1.89.0)
• 結果: AC
• 実行時間: 1,519 ms / 10,000 ms
• コード長: 9,522 bytes
• コンパイル時間: 1,852 ms
• コンパイル使用メモリ: 103,416 KB
• 最終ジャッジ日時: 2025-02-07 23:53:33
• ジャッジサーバーID （参考情報）: judge2 / judge3

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

ソースコード

#line 1 "Main.cpp"
#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


#line 2 "nachia\\math\\modulo-primitive-root.hpp"

#line 5 "nachia\\math\\modulo-primitive-root.hpp"

namespace nachia{

template<unsigned int MOD>
struct PrimitiveRoot{
    static constexpr unsigned long long powm(unsigned long long a, unsigned long long i) {
        unsigned long long res = 1, aa = a;
        while(i){
            if(i & 1) res = res * aa % MOD;
            aa = aa * aa % MOD;
            i /= 2;
        }
        return res;
    }
    static constexpr bool examine_val(unsigned int g){
        unsigned int t = MOD - 1;
        for(unsigned long long d=2; d*d<=t; d++) if(t % d == 0){
            if(powm(g, (MOD - 1) / d) == 1) return false;
            while(t % d == 0) t /= d;
        }
        if(t != 1) if(powm(g, (MOD - 1) / t) == 1) return false;
        return true;
    }
    static constexpr unsigned int get_val(){
        for(unsigned int x=2; x<MOD; x++) if(examine_val(x)) return x;
        return 0;
    }
    static const unsigned int val = get_val();
};

}
#line 127 "Main.cpp"

#include <iostream>
#include <algorithm>
#include <atcoder/modint>
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++)
using Modint = atcoder::static_modint<998244353>;

 

void NTT(vector<Modint>& A){
    Modint g = nachia::PrimitiveRoot<Modint::mod()>::get_val();
    constexpr unsigned int MOD = Modint::mod();
    int N = 1;
    while (N < (int)A.size()) N *= 2;
    static vector<Modint> exp_i_revbit_diff = { g.pow((MOD - 1) >> 2) };
    for(int i=exp_i_revbit_diff.size(); (1<<(i+1))<N; i++){
        Modint diffdiff = g.pow((MOD - 1) / 2 + ((MOD - 1) >> (i+2)) * 3);
        exp_i_revbit_diff.push_back(diffdiff);
    }
    for (int i = 1; i < N; i <<= 1) {
        Modint 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++) {
                Modint l = A[k];
                Modint r = A[k + maxk] * q;
                A[k] = l + r;
                A[k + maxk] = l - r;
            }
            if(j+1 != i) for(int d=0; ; d++) if(!(j&(1<<d))){
                q = q * exp_i_revbit_diff[d];
                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 u32 MOD = 998244353;
const u32 NTTzeta = 3;


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

Modint k0;
vector<Modint> inv_mod;
vector<Modint> C;
vector<vector<Modint>> nttC;
vector<Modint> inv_ntt_size;

vector<Modint> 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<Modint> sigma(const vector<Modint>& B){
    int n = B.size() + 2;
    if(n >= 10){
      int Z = 1, d = 0;
      while(Z < n){ Z *= 2; d++; }
      vector<Modint> revB(Z*2, 0);
      rep(i, B.size()) revB[Z-i] = B[i];
      NTT(revB);
      rep(i, Z*2) revB[i] = revB[i] * nttC[d+1][i];
      NTT(revB);
      reverse(revB.begin() + 1, revB.end());
      rep(i, n) revB[i] = revB[i+Z] * inv_ntt_size[d+1];
      revB.resize(n);
      return revB;
    }
    vector<Modint> ans(n);
    for(int i=0; i<n; i++){
      for(int j=0; j<(int)B.size(); j++){
        ans[i] += C[i+j] * B[j];
      }
    }
    return ans;
}

vector<Modint> sigma_tree(const nachia::CentroidDecomposition::BFSUnit& tree){
    vector<Modint> dist_freq(tree.I.size(), 0);
    get_bfsbuf_dist(tree);
    for(int p : tree.I) dist_freq[bfsbuf_dist[p]] += Q[p];
    while(1 < dist_freq.size() && dist_freq.back().val() == 0) dist_freq.pop_back();
    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 *= i;
    k0 *= k0;
    inv_mod.assign(max_ntt_size+1, 1);
    for(int i=2; i<=max_ntt_size; i++) inv_mod[i] = (MOD / i) * -inv_mod[MOD % i];
    C.resize(max_ntt_size);
    for(int i=0; i<max_ntt_size; i++) C[i] = k0 * inv_mod[i+1] * inv_mod[i+1];

    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];
        NTT(nttC[d]);
        inv_ntt_size[d] = Modint(1<<d).inv();
    }

    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<Modint> 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<Modint> sigma_nx = sigma_tree(bfs_nx);
            for(int p : bfs_nx.I){
                int d = bfsbuf_dist[p] + 1;
                ans[p] += sigma_s[d] - sigma_nx[d+1];
            }
        }

        ans[s] += sigma_s[0];
    }

    rep(p,N) cout << ans[p].val() << '\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;