kmjp's blog

競技プログラミング参加記です

AtCoder ABC #163 : F - path pass i

なんか妙に手間取った。
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を繰り返した。
サンプルが弱めだったのかな。