毎日格上の問題を倒すやつの93日目です。
問題リンク
問題概要
頂点に色がついた木があるよ。
各色について、その色を一回以上通るパスの個数を求めてね。
解法
頂点0を根とした根つき木とします。
トポロジカルソートをして、オイラーツアー順を取得しておきます。これはどちらもDFSで対応できます。
各色、ソート順が大きい方から、「その頂点を通って、以前に処理した頂点たちは通らないパスの総数」を求めたいです。
これは、先ほど作ったオイラーツアーから遅延セグメント木を作成すれば、部分木の削除と部分木の頂点数の取得を高速にできます。
よって、これらからいい感じに計算すればよいです。
提出コード
#include<bits/stdc++.h> #include<atcoder/lazysegtree> using namespace std; using namespace atcoder; using ll = long long; #define rep(i,n) for(int i=0; i<n; i++) #pragma GCC target("avx2") #pragma GCC optimize("O3") #pragma GCC optimize("unroll-loops") void dfs(int pos,int bef,int &cnt, vector<int> &c,vector<vector<int>> &tree,vector<vector<int>> &col,vector<pair<int,int>> &tour){ tour[pos].first=cnt; cnt++; col[c[pos]-1].push_back(pos); for(int i:tree[pos]){ if(i!=bef){ dfs(i,pos,cnt,c,tree,col,tour); } } tour[pos].second=cnt; cnt++; } pair<int,int> op(pair<int,int> x,pair<int,int> y){ return pair<int,int>(x.first+y.first,x.second+y.second); } pair<int,int> e(){ return pair<int,int>(0,0); } pair<int,int> mapp(int f,pair<int,int> x){ if(f==-1)return x; return pair<int,int>(x.second*f,x.second); } int comp(int f,int g){ if(f==-1)return g; return f; } int id(){ return -1; } int main(){ ios::sync_with_stdio(false); cin.tie(nullptr); int N; cin>>N; vector<int> c(N); rep(i,N){ cin>>c[i]; } vector<vector<int>> tree(N); rep(i,N-1){ int a,b; cin>>a>>b; tree[a-1].push_back(b-1); tree[b-1].push_back(a-1); } vector<vector<int>> col(N); vector<pair<int,int>> tour(N); int cnt=0; dfs(0,-1,cnt,c,tree,col,tour); lazy_segtree<pair<int,int>,op,e,int,mapp,comp,id> seg(vector<pair<int,int>>(N*2,pair<int,int>(1,1))); rep(i,N){ seg.apply(0,N*2,1); reverse(col[i].begin(),col[i].end()); ll ans=0; for(int j:col[i]){ ll add=0,squ=0; for(int k:tree[j]){ if(tour[j].first<tour[k].first){ ll size=seg.prod(tour[k].first,tour[k].second+1).first/2; add+=size; squ+=size*size; } } ll size=(seg.all_prod().first-seg.prod(tour[j].first,tour[j].second+1).first)/2; add+=size; squ+=size*size; ans+=(add*add-squ)/2+seg.all_prod().first/2; seg.apply(tour[j].first,tour[j].second+1,0); } cout<<ans<<endl; } }
感想
オイラーツアーだけでなくいろんな典型が使えるいい問題でした
