なんか妙に手間取った。
https://atcoder.jp/contests/abc163/tasks/abc163_f
問題
木を成す無向グラフが与えられる。
各頂点vには色C[v]が塗られている。
この木上のパスのうち、ある色の頂点の頂点を1個以上通るものについて、色ごとに数を答えよ。
解法
Editorialとは異なり、それぞれ数え上げた。
まず、オイラーツアーで各頂点の親子関係を求めたうえ、BITで各頂点自身の親方向の頂点数を求めておく。
以後、各色について、その色を持つ頂点をオイラーツアー順に処理していく。
今見ている頂点を通るパスの数は、異なる子頂点または親頂点側に属する2点を選んだ時である。
オイラーツアーを先に行っているので、子頂点側のSubtreeの頂点数は容易にわかる。
親側の頂点数は、先のBITでわかる。
頂点vを処理したら、その頂点vの切り離しを行う。すなわち、頂点の親側の各頂点には、substee分の頂点を引き、子頂点のsubtreeには、親頂点及び他の子頂点のsubtree側の頂点数を引く。
これらはBITで区間加算・区間減算を行うことで対応できる。
一点注意点として、vを切り離すとき、vの親側の頂点で影響を受ける範囲は、vの祖先の最寄りの同色の頂点をwとするとwより1つv側の頂点のsubtreeのうちvのsubtreeを除いた範囲である。
そこでこのような点はオイラーツアーのついでに求めておく。
int N; int C[202020]; vector<pair<int,int>> cand[202020]; vector<int> E[202020]; int L[202020],R[202020],id; int rem[201010]; int pr[202020]; vector<int> hist[202020]; template<class V, int ME> class BIT { public: V bit[1<<ME]; V operator()(int e) {if(e<0) return 0;V s=0;e++;while(e) s+=bit[e-1],e-=e&-e; return s;} void add(int e,V v) { e++; while(e<=1<<ME) bit[e-1]+=v,e+=e&-e;} }; BIT<ll,20> bt; void dfs(int cur,int pre) { if(hist[C[cur]].size()) pr[cur]=hist[C[cur]].back(); L[cur]=id++; cand[C[cur]].push_back({L[cur],cur}); FORR(e,E[cur]) { hist[C[cur]].push_back(e); if(e!=pre) dfs(e,cur); hist[C[cur]].pop_back(); } R[cur]=id; bt.add(L[cur],N-(R[cur]-L[cur])); bt.add(L[cur]+1,-(N-(R[cur]-L[cur]))); } void solve() { int i,j,k,l,r,x,y; string s; cin>>N; FOR(i,N) { cin>>C[i]; C[i]--; } FOR(i,N-1) { cin>>x>>y; E[x-1].push_back(y-1); E[y-1].push_back(x-1); } dfs(0,0); FOR(i,N) { ll ret=0; sort(ALL(cand[i])); FORR(c,cand[i]) { x=c.second; int num=bt(L[x])+1; ret+=num; FORR(e,E[x]) if(L[e]>L[x]) { ret+=1LL*num*(R[e]-L[e]); num+=(R[e]-L[e]); } //cout<<x<<" "<<ret<<endl; FORR(e,E[x]) if(L[e]>L[x]) { bt.add(L[e],-(num-(R[e]-L[e]))); bt.add(R[e],num-(R[e]-L[e])); } y=pr[x]; bt.add(L[y],-(R[x]-L[x])); bt.add(L[x],(R[x]-L[x])); bt.add(R[x],-(R[x]-L[x])); bt.add(R[y],(R[x]-L[x])); rem[x]=num; } FORR(c,cand[i]) { x=c.second; FORR(e,E[x]) if(L[e]>L[x]) { bt.add(L[e],(rem[x]-(R[e]-L[e]))); bt.add(R[e],-(rem[x]-(R[e]-L[e]))); } y=pr[x]; bt.add(L[y],(R[x]-L[x])); bt.add(L[x],-(R[x]-L[x])); bt.add(R[x],(R[x]-L[x])); bt.add(R[y],-(R[x]-L[x])); } cout<<ret<<endl; } }
まとめ
足し引きする範囲を間違えてWAを繰り返した。
サンプルが弱めだったのかな。