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

提出ソース

結果

問題	No.1796 木上のクーロン
ユーザー	👑 Nachia
提出日時	2021-12-13 22:36:55
言語	C++17 (gcc 12.3.0 + boost 1.83.0)
結果	AC
実行時間	2,230 ms / 10,000 ms
コード長	6,801 bytes
コンパイル時間	1,471 ms
コンパイル使用メモリ	102,412 KB
実行使用メモリ	51,256 KB
最終ジャッジ日時	2024-09-19 22:45:06
合計ジャッジ時間	25,513 ms
ジャッジサーバーID （参考情報）	judge3 / judge5

このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	2 ms 6,816 KB
testcase_01	AC	2 ms 6,944 KB
testcase_02	AC	2 ms 6,940 KB
testcase_03	AC	1 ms 6,944 KB
testcase_04	AC	0 ms 6,940 KB
testcase_05	AC	1 ms 6,940 KB
testcase_06	AC	2 ms 6,944 KB
testcase_07	AC	2 ms 6,940 KB
testcase_08	AC	3 ms 6,940 KB
testcase_09	AC	3 ms 6,944 KB
testcase_10	AC	3 ms 6,944 KB
testcase_11	AC	3 ms 6,944 KB
testcase_12	AC	3 ms 6,940 KB
testcase_13	AC	2 ms 6,940 KB
testcase_14	AC	3 ms 6,940 KB
testcase_15	AC	4 ms 6,944 KB
testcase_16	AC	4 ms 6,944 KB
testcase_17	AC	2 ms 6,940 KB
testcase_18	AC	3 ms 6,944 KB
testcase_19	AC	4 ms 6,940 KB
testcase_20	AC	295 ms 15,192 KB
testcase_21	AC	266 ms 15,324 KB
testcase_22	AC	622 ms 27,224 KB
testcase_23	AC	604 ms 27,220 KB
testcase_24	AC	984 ms 42,064 KB
testcase_25	AC	1,010 ms 43,980 KB
testcase_26	AC	1,418 ms 50,768 KB
testcase_27	AC	1,414 ms 51,144 KB
testcase_28	AC	2,230 ms 50,876 KB
testcase_29	AC	2,214 ms 50,876 KB
testcase_30	AC	639 ms 50,404 KB
testcase_31	AC	665 ms 50,244 KB
testcase_32	AC	2,224 ms 50,876 KB
testcase_33	AC	1,411 ms 51,256 KB
testcase_34	AC	1,357 ms 51,224 KB
testcase_35	AC	2,017 ms 50,880 KB
testcase_36	AC	2,019 ms 50,628 KB

権限があれば一括ダウンロードができます

ソースコード

raw source code

#include <iostream>
#include <vector>
#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<(n); i++)

/////////////////////////////////////
// Problem : https://yukicoder.me/problems/7417
// Time Complexity : O( N log^2 N )

/////////////////////////////////////
// input
//
// N M K
// Q[1] Q[2] ... Q[N]
// u[1] v[1]
// u[2] v[2]
//  :
// u[N-1] v[N-1]
//


struct CentroidDecomposition {
    int n;
    vector<vector<int>> E;
    vector<int> cdep;
    vector<int> cp;
    vector<int> cbfs;
    int maxdep;
    CentroidDecomposition(vector<vector<int>> edges) {
        E = move(edges);
        n = E.size();
        vector<int> Z(n, 1);
        {
            vector<int> P(n, -1);
            vector<int> I = { 0 };
            rep(i, I.size()) {
                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);
        vector<pair<int, int>> I = { {0,-1} };
        rep(i, I.size()) {
            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(make_pair(e, s));
            }
        }
        maxdep = 0;
        for (int a : cdep) maxdep = max(maxdep, a);
    }

    struct BFSUnit {
        vector<int> I;
        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);
        rep(i, res.I.size()) {
            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;
    }
};



// 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 < 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;
vector<vector<int>> 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(CentroidDecomposition::BFSUnit tree){
    bfsbuf_dist[tree.I.front()] = 0;
    for(int i=1; i<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 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];
    E.resize(N);
    rep(i,N-1){
        int u,v; cin >> u >> v; u--; v--;
        E[u].push_back(v);
        E[v].push_back(u);
    }
    
    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 = 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){
            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;
}