結果
問題 | No.2115 Making Forest Easy |
ユーザー |
![]() |
提出日時 | 2022-10-28 22:41:20 |
言語 | C++17 (gcc 13.3.0 + boost 1.87.0) |
結果 |
AC
|
実行時間 | 782 ms / 2,000 ms |
コード長 | 2,719 bytes |
コンパイル時間 | 4,914 ms |
コンパイル使用メモリ | 262,484 KB |
最終ジャッジ日時 | 2025-02-08 14:49:50 |
ジャッジサーバーID (参考情報) |
judge2 / judge1 |
(要ログイン)
ファイルパターン | 結果 |
---|---|
sample | AC * 2 |
other | AC * 50 |
ソースコード
#include <stdio.h> #include <atcoder/all> #include <bits/stdc++.h> using namespace std; using namespace atcoder; using mint = modint998244353; #define rep(i,n) for (int i = 0; i < (n); ++i) #define Inf32 1000000001 #define Inf64 4000000000000000001 template <class Ss, Ss (*op0)(Ss, Ss), Ss (*e)(), class Sv, Ss (*op1)(Sv, Ss), class Se, Ss (*op2)(Se, Ss)> struct rerooting{ int n; vector<vector<int>> to; vector<vector<Se>> edge; vector<Ss> subtree,answer; vector<Sv> vertex; rerooting(int _n){ n = _n; to.resize(n); edge.resize(n); subtree.resize(n,e()); answer.resize(n,e()); vertex.resize(n); } rerooting(vector<Sv> v){ n = v.size(); to.resize(n); edge.resize(n); subtree.resize(n,e()); answer.resize(n,e()); vertex = v; } void add_directed_edge(int u,int v,Se x){ to[u].push_back(v); edge[u].push_back(x); } void add_edge(int u,int v,Se x){ add_directed_edge(u,v,x); add_directed_edge(v,u,x); } void solve(){ dfs(0,-1); redfs(0,-1,e()); } void dfs(int cur,int par){ Ss temp = e(); for(int i=0;i<to[cur].size();i++){ int nxt = to[cur][i]; if(nxt==par)continue; dfs(nxt,cur); temp = op0(temp, op2(edge[cur][i], subtree[nxt])); } subtree[cur] = op1(vertex[cur], temp); } void redfs(int cur, int par, Ss ps){ vector<Ss> P(to[cur].size()+1,e()); rep(i,to[cur].size()){ int nxt = to[cur][i]; if(nxt!=par)P[i+1] = op2(edge[cur][i], subtree[nxt]); else P[i+1] = op2(edge[cur][i], ps); } vector<Ss> S(P.rbegin(),P.rend()-1); S.insert(S.begin(),e()); rep(i,to[cur].size()){ P[i+1] = op0(P[i], P[i+1]); S[i+1] = op0(S[i+1], S[i]); } answer[cur] = op1(vertex[cur], P.back()); for(int i=0;i<to[cur].size();i++){ int nxt = to[cur][i]; if(nxt!=par){ redfs(nxt,cur,op1(vertex[cur], op0(P[i],S[to[cur].size()-1-i]))); } } } }; using A = array<mint,1005>; A e(){ A ret; rep(i,1005)ret[i] = 0; ret[0] = 1; return ret; } A op0(A a,A b){ rep(i,1004){ a[i+1] += a[i]; b[i+1] += b[i]; } A ret = e(); ret[0] = 0; rep(i,1005){ if(i==0)ret[i] = a[i]*b[i]; else{ ret[i] += (a[i]-a[i-1]) * (b[i]-b[i-1]); ret[i] += (a[i]-a[i-1]) * b[i-1]; ret[i] += a[i-1] * (b[i]-b[i-1]); } } return ret; } A op1(int a,A b){ rep(i,a){ b[a] += b[i]; b[i] = 0; } return b; } A op2(int a,A b){ mint s = 0; rep(i,1003)s += b[i]; b[0] += s; return b; } int main(){ int n; cin>>n; vector<int> a(n); rep(i,n)cin>>a[i]; rerooting<A,op0,e,int,op1,int,op2> R(a); rep(i,n-1){ int u,v; cin>>u>>v; R.add_edge(u-1,v-1,1); } R.solve(); mint ans = 0; rep(i,n){ rep(j,1003){ ans += R.answer[i][j] * j; } } cout<<ans.val()<<endl; return 0; }