結果

問題 No.1796 木上のクーロン
コンテスト
ユーザー Nachia
提出日時 2021-12-19 01:08:15
言語 Java
(openjdk 23)
結果
AC  
実行時間 5,768 ms / 10,000 ms
コード長 5,234 bytes
コンパイル時間 2,164 ms
コンパイル使用メモリ 78,152 KB
実行使用メモリ 119,860 KB
最終ジャッジ日時 2024-09-19 22:50:21
合計ジャッジ時間 60,728 ms
ジャッジサーバーID
（参考情報） 		judge3 / judge2
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 34
権限があれば一括ダウンロードができます

ソースコード

diff # 

import java.io.PrintWriter;
import java.util.Scanner;

public class Main {

	static final int MOD = 998244353;
	static final long NTTg = 3;

	static long powm(long a, long i) {
		if(i == 0) return 1;
		long r = powm(a*a%MOD,i/2);
		if(i%2 == 1) r=r*a%MOD;
		return r;
	}

	static final long invNTTg = powm(NTTg, MOD-2);

	static void NTT(long[] A, long g) {
		int N = A.length;
		for(int i=0, j=0; j<N; j++) {
			if(i < j) {
				long b = A[i];
				A[i] = A[j];
				A[j] = b;
			}
			for(int k=N>>1; k>(i^=k); k>>=1);
		}
		for(int i=1; i<N; i<<=1) {
			long q = powm(g,(MOD-1)/i/2), qj = 1;
			for(int j=0; j<i; j++) {
				for(int k=j; k<N; k+=i*2) {
					long l = A[k];
					long r = A[k+i] * qj % MOD;
					A[k] = l+r;
					if(A[k] >= MOD) A[k] -= MOD;
					A[k+i] = l+MOD-r;
					if(A[k+i] >= MOD) A[k+i] -= MOD;
				}
				qj = qj * q % MOD;
			}
		}
	}

	int N;
	long Q[];
	int E[][];

	int cdep[];
	int cp[];
	int cbfs[];

	void centroid_decomposition() {
		cdep = new int[N];
		cp = new int[N]; for(int i=0; i<N; i++) cp[i] = -1;
		cbfs = new int[N];
		int P[] = new int[N]; for(int i=0; i<N; i++) P[i] = -1;
		int I[] = new int[N];
		int Iidx = 0; I[Iidx++] = 0;
		for(int p : I) for(int e : E[p]) if(P[p] != e) { P[e] = p; I[Iidx++] = e; }
		int Z[] = new int[N]; for(int i=0; i<N; i++) Z[i] = 1;
		for(int i=N-1; i>=1; i--) Z[P[I[i]]] += Z[I[i]];

		int cdP[] = new int[N];
		Iidx = 0;
		cdP[Iidx] = -1;
		I[Iidx++] = 0;
		for(int i=0; i<N; i++) {
			int s = I[i];
			int par = cdP[i];
			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[i] = s;
			Z[s] = 0;
			if(par != -1) { cdep[s] = cdep[par]+1; cp[s] = par; }
			for(int e : E[s]) if(Z[e] != 0) {
				cdP[Iidx] = s;
				I[Iidx++] = e;
			}
		}
	}



	long k0;
	int max_ntt_size_log;
	int max_ntt_size;
	long inv_mod[];
	long C[];
	long nttC[][];

	void input() {
		Scanner sc = new Scanner(System.in);
		N = sc.nextInt();
		Q = new long[N];
		for(int i=0; i<N; i++) Q[i] = (long)sc.nextInt();
		{
			int edge_idx[] = new int[N];
			int edge_u[] = new int[N-1];
			int edge_v[] = new int[N-1];
			for(int i=0; i<N-1; i++) {
				int u = sc.nextInt() - 1;
				int v = sc.nextInt() - 1;
				edge_idx[u]++;
				edge_idx[v]++;
				edge_u[i] = u;
				edge_v[i] = v;
			}
			E = new int[N][];
			for(int i=0; i<N; i++) E[i] = new int[edge_idx[i]];
			for(int i=0; i<N-1; i++) {
				int u = edge_u[i];
				int v = edge_v[i];
				E[u][--edge_idx[u]] = v;
				E[v][--edge_idx[v]] = u;
			}
		}
		sc.close();
	}

	void init_modint() {
		max_ntt_size_log = 0;
		while((1 << max_ntt_size_log) < N+6) max_ntt_size_log++; max_ntt_size_log++;
		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 = new long[max_ntt_size+1];
		inv_mod[1] = 1;
		for(int i=2; i<=max_ntt_size; i++) inv_mod[i] = MOD - MOD / i * inv_mod[MOD%i] % MOD;

		C = new long[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 = new long[max_ntt_size_log+1][];
		for(int d=0; d<=max_ntt_size_log; d++) {
			long inv_ntt_size = powm(1<<d, MOD-2);
			nttC[d] = new long[1<<d];
			for(int i=0; i<1<<d; i++) nttC[d][i] = C[i] * inv_ntt_size % MOD;
			NTT(nttC[d], NTTg);
		}
	}

	int bfsbuf_dist[];
	int bfsbuf_parent[];
	int bfsbuf_I[];
	int bfsbuf_I_size;
	long[] sigma_tree(int s, int dep) {
		if(cdep[s] < dep) return new long[1];
		bfsbuf_dist[s] = 0;
		bfsbuf_parent[s] = -1;
		bfsbuf_I_size = 0;
		bfsbuf_I[bfsbuf_I_size++] = s;
		int maxdist = 0;

		for(int i=0; i<bfsbuf_I_size; i++) {
			int p = bfsbuf_I[i];
			maxdist = Math.max(maxdist, bfsbuf_dist[p]);
			for(int e : E[p]) if(bfsbuf_parent[p] != e) {
				if(cdep[e] < dep) continue;
				bfsbuf_parent[e] = p;
				bfsbuf_dist[e] = bfsbuf_dist[p] + 1;
				bfsbuf_I[bfsbuf_I_size++] = e;
			}
		}

		long dfreq[] = new long[maxdist+1];
		for(int i=0; i<bfsbuf_I_size; i++) {
			int p = bfsbuf_I[i];
			int d = bfsbuf_dist[p];
			dfreq[d] = (dfreq[d] + Q[p]) % MOD;
		}

		int Z = 1;
		int d = 0;
		while(Z < dfreq.length+2) { Z*=2; d++; }
		long res[] = new long[Z*2];
		for(int i=0; i<dfreq.length; i++) res[Z-i] = dfreq[i];
		NTT(res, NTTg);
		for(int i=0; i<Z*2; i++) res[i] = res[i] * nttC[d+1][i] % MOD;
		NTT(res, invNTTg);
		for(int i=0; i<Z; i++) res[i] = res[i+Z];
		return res;
	}

	void solve() {
		input();
		init_modint();
		centroid_decomposition();

		bfsbuf_dist = new int[N];
		bfsbuf_parent = new int[N];
		bfsbuf_I = new int[N];

		long ans[] = new long[N];
		for(int s=0; s<N; s++) {
			int dep_s = cdep[s];
			long[] sigma_s = sigma_tree(s, dep_s);
			for(int nx : E[s]) if(cdep[nx] > dep_s) {
				long[] sigma_nx = sigma_tree(nx, dep_s + 1);
				for(int i=0; i<bfsbuf_I_size; i++) {
					int p = bfsbuf_I[i];
					int d = bfsbuf_dist[p] + 1;
					ans[p] += sigma_s[d] - sigma_nx[d+1] + MOD;
				}
			}
			ans[s] += sigma_s[0];
		}
		for(int i=0; i<N; i++) ans[i] %= MOD;

		PrintWriter wt = new PrintWriter(System.out);
		for(int i=0; i<N; i++) {
			wt.println(ans[i]);
		}
		wt.flush();
	}



	public static void main(String[] args) {
		Main solver = new Main();
		solver.solve();
	}
}
0