結果
問題 | No.2115 Making Forest Easy |
ユーザー |
|
提出日時 | 2022-11-06 23:02:25 |
言語 | C++14 (gcc 13.3.0 + boost 1.87.0) |
結果 |
AC
|
実行時間 | 376 ms / 2,000 ms |
コード長 | 5,360 bytes |
コンパイル時間 | 3,593 ms |
コンパイル使用メモリ | 236,760 KB |
実行使用メモリ | 45,056 KB |
最終ジャッジ日時 | 2024-07-20 09:57:18 |
合計ジャッジ時間 | 18,200 ms |
ジャッジサーバーID (参考情報) |
judge3 / judge1 |
(要ログイン)
ファイルパターン | 結果 |
---|---|
sample | AC * 2 |
other | AC * 50 |
ソースコード
#include<bits/stdc++.h> #include<atcoder/all> #define rep(i,b) for(int i=0;i<b;i++) #define rrep(i,b) for(int i=b-1;i>=0;i--) #define rep1(i,b) for(int i=1;i<b;i++) #define repx(i,x,b) for(int i=x;i<b;i++) #define rrepx(i,x,b) for(int i=b-1;i>=x;i--) #define fore(i,a) for(auto i:a) #define fore1(i,a) for(auto &i:a) #define rng(x) (x).begin(), (x).end() #define rrng(x) (x).rbegin(), (x).rend() #define sz(x) ((int)(x).size()) #define pb push_back #define fi first #define se second #define pcnt __builtin_popcountll using namespace std; using namespace atcoder; using ll = long long; using ld = long double; template<typename T> using mpq = priority_queue<T, vector<T>, greater<T>>; template<typename T> bool chmax(T &a, const T &b) { if (a<b) { a=b; return 1; } return 0; } template<typename T> bool chmin(T &a, const T &b) { if (b<a) { a=b; return 1; } return 0; } template<typename T> ll sumv(const vector<T>&a){ll res(0);for(auto&&x:a)res+=x;return res;} bool yn(bool a) { if(a) {cout << "Yes" << endl; return true;} else {cout << "No" << endl; return false;}} #define dame { cout << "No" << endl; return;} #define dame1 { cout << -1 << endl; return;} #define deb(x,y) cout << x << " " << y << endl; #define debp(p) cout << p.fi << " " << p.se << endl; #define out cout << ans << endl; #define outd cout << fixed << setprecision(20) << ans << endl; #define outm cout << ans.val() << endl; #define outv fore(yans , ans) cout << yans << "\n"; #define outdv fore(yans , ans) cout << yans.val() << "\n"; #define showv(v) {fore(vy , v) {cout << vy << " ";} cout << endl;} #define showv2(v) fore(vy , v) cout << vy << "\n"; #define showvm(v) {fore(vy , v) {cout << vy.val() << " ";} cout << endl;} #define showvm2(v) fore(vy , v) cout << vy.val() << "\n"; using pll = pair<ll,ll>;using pil = pair<int,ll>;using pli = pair<ll,int>;using pii = pair<int,int>;using pdd = pair<ld,ld>; using tp = tuple<int ,int ,int>; using vi = vector<int>;using vd = vector<ld>;using vl = vector<ll>;using vs = vector<string>;using vb = vector<bool>; using vpii = vector<pii>;using vpli = vector<pli>;using vpll = vector<pll>;using vpil = vector<pil>; using vvi = vector<vector<int>>;using vvl = vector<vector<ll>>;using vvs = vector<vector<string>>;using vvb = vector<vector<bool>>; using vvpii = vector<vector<pii>>;using vvpli = vector<vector<pli>>;using vvpll = vector<vpll>;using vvpil = vector<vpil>; using mint = modint998244353; using vm = vector<mint>; using vvm = vector<vector<mint>>; vector<int> dx={1,0,-1,0,1,1,-1,-1},dy={0,1,0,-1,1,-1,1,-1}; ll gcd(ll a, ll b) { return a?gcd(b%a,a):b;} ll lcm(ll a, ll b) { return a/gcd(a,b)*b;} #define yes {cout <<"Yes"<<endl;} #define yesr { cout <<"Yes"<<endl; return;} #define no {cout <<"No"<<endl;} #define nor { cout <<"No"<<endl; return;} const double eps = 1e-10; const ll LINF = 1001002003004005006ll; const int INF = 1001001001; #ifdef MY_LOCAL_DEBUG #define show(x) cerr<<#x<<" = "<<x<<endl #define showp(p) cerr<<#p<<" = "<<p.fi<<" : "<<p.se<<endl #define show2(x,y) cerr<<#x<<" = "<<x<<" : "<<#y<<" = "<<y<<endl #define show3(x,y,z) cerr<<#x<<" = "<<x<<" : "<<#y<<" = "<<y<<" : "<<#z<<" = "<<z<<endl #define show4(x,y,z,x2) cerr<<#x<<" = "<<x<<" : "<<#y<<" = "<<y<<" : "<<#z<<" = "<<z<<" : "<<#x2<<" = "<<x2<<endl #define test(x) cout << "test" << x << endl #else #define show(x) #define showp(p) #define show2(x,y) #define show3(x,y,z) #define show4(x,y,z,x2) #define test(x) #endif void solve(){ int n; cin>>n; vi a(n); rep(i,n) cin>>a[i]; vvi e(n); rep(i,n-1){ int a,b; cin>>a>>b; a--; b--; e[a].pb(b); e[b].pb(a); } const int mxa = 1001; vvm dp1(n , vm(mxa)),dp2(n , vm(mxa)); vi dd(n); auto dfs = [&](auto dfs,int x,int pre)->int{ int d = 1; dp1[x][a[x]] = 1; dp2[x][a[x]] = 1; fore(y , e[x]){ if(y == pre) continue; int k = dfs(dfs , y ,x); vm tmp1(mxa); vm tmp2(mxa); swap(tmp1 , dp1[x]); swap(tmp2 , dp2[x]); vm sumx1(mxa),sumy1(mxa),sumx2(mxa),sumy2(mxa); dp1[y][0] = mint(2).pow(k-1); dp2[y][0] = 0; rep(j,mxa){ sumx1[j] = tmp1[j]; sumy1[j] = dp1[y][j]; sumx2[j] = tmp2[j]; sumy2[j] = dp2[y][j]; if (j != 0){ sumx1[j] += sumx1[j-1]; sumy1[j] += sumy1[j-1]; sumx2[j] += sumx2[j-1]; sumy2[j] += sumy2[j-1]; } } rep(j,mxa){ dp1[x][j] = sumx1[j]*dp1[y][j] + tmp1[j]*sumy1[j] - tmp1[j]*dp1[y][j]; dp2[x][j] += sumx2[j]*dp1[y][j] + tmp1[j]*sumy2[j]; dp2[x][j] += sumx1[j]*dp2[y][j] + tmp2[j]*sumy1[j]; dp2[x][j] -= tmp1[j]*dp2[y][j] + tmp2[j]*dp1[y][j]; } d += k; } dd[x] = d; return d; }; dfs(dfs,0,-1); mint ans = 0; rep(i,n) rep1(j,mxa){ int ik = 0; if (i == 0) ik = 1; ans += dp2[i][j] * j * mint(2).pow(n-1-dd[i]+ik); } outm; return; } int main(){ ios::sync_with_stdio(false); cin.tie(0); int t = 1; //cin>>t; rep(i,t){ solve(); } return 0; }