kmjp's blog

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

Codeforces #221 Div1. D. Tree and Queries

Codeforces

本番は「計算量が落とせないな…たぶん平方分割なんだろうけど…」と思っているうちに時間切れ。
後日「サブツリーの点の合計の情報をマージしていけばいいのかな？」と思いついた。
前者は他人の回答を参考にして、後者は結局自力で2種類の解き方ができた。
http://codeforces.com/contest/375/problem/D

問題

N点からなる根つき木を成すグラフが与えられる。
各頂点には色C[i]が塗られている。

ここで、配列V,Kで与えられる以下のM個のクエリを処理せよ。

  • 点V[i]に対し、同じ色の点がK[i]個以上あるような色の数を答えよ。

解法

V[i]が木の根に近い方の点を指す場合、サブツリーに含まれる点の色を数え上げるのにO(N)かかるため、全体でO(NM)かかる。
なんとかしてこれを軽くしたい。

平方分割法と、サブツリーの情報をマージしていく方法とどちらでも解ける。
前提として、事前にDFSでEuler Tourを求めておき、各頂点を1列に並べたうえで、各頂点V[i]に対するサブツリーの範囲L[i]～R[i]を求めておく。

１：平方分割

各クエリの点V[i]に対応するサブツリー範囲(L[i],R[i])について、以下のようにある分割数Sに対し、各クエリをソートする。

  • L[i]/S == L[j]/S ならR[i]とR[j]の大小でソートする。
  • L[i]/S != L[j]/S ならL[i]とL[j]の大小でソートする。

このようにソートすると、i番目のクエリの結果に対しi+1番目を次のように求められる。

  • L[i]＜L[i+1]ならL[i]～(L[i+1]-1)番の頂点の色の分、色の頂点数をデクリメントする
  • L[i]＞L[i+1]ならL[i+1]～(L[i]-1)番の頂点の色の分、色の頂点数をインクリメントする
  • R[i]＜R[i+1]なら(R[i]+1)～R[i+1]番の頂点の色の分、色の頂点数をインクリメントする
  • R[i]＞R[i+1]なら(R[i+1]+1)～R[i]番の頂点の色の分、色の頂点数をデクリメントする

L[i]/S==L[i+1]/Sなら、L[i]～L[i+1]の移動距離は高々Sである。
また、L[i]/S==L[i+1]/SならR[i]<=R[i+1]なので、L[i]/S==L[i+a]/Sであるようなaについて、R[i]～R[i+a]は高々移動距離はNである。

クエリの各頂点がおおむね均等に分配しているとすると、L[i]～L[i+1]の移動距離がS程度なのでMクエリでO(SM)、R[0]～R[M-1]の移動距離は、N/S回ほどR[i]が0～(N-1)の間を巡回するのでO(N/S*M)。

よって両者でO(M(N/S + S))なので、S=√N程度にすると全体をO(M√N)で処理できる。

int N,M;
int ti=0;
int C[100010],num[100010];
int L[100010],R[100010],C2[100010];
vector<int> E[100010];
int Q[100010],LQ[100010],RQ[100010],K[100010],ID[100010];
int cs[100010],csc[100010],res[100010];

void dfs(int cur,int pre) {
	int i;
	int pi=-1;
	
	L[cur]=ti++;
	C2[L[cur]]=C[cur];
	num[cur]=1;
	FOR(i,E[cur].size()) {
		int tar=E[cur][i];
		if(pre!=tar) dfs(tar,cur), num[cur]+=num[tar];
	}
	R[cur]=ti-1;
}

int cmp(int a,int b) {
	if(LQ[a]==LQ[b]) return RQ[a]<RQ[b];
	return LQ[a]<LQ[b];
}

void solve() {
	int f,i,j,k,l,x,y;
	
	cin>>N>>M;
	FOR(i,N) cin>>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,-1);
	FOR(i,M) {
		cin>>x>>K[i];
		Q[i]=x-1;
		ID[i]=i;
		LQ[i]=L[Q[i]]/333;
		RQ[i]=R[Q[i]];
	}
	sort(ID,ID+M,cmp);
	int cl=0,cr=0;
	cs[C2[0]]=1;
	csc[1]=1;
	FOR(i,M) {
		j=Q[ID[i]];
		while(cl<L[j]) csc[cs[C2[cl++]]--]--;
		while(cr<R[j]) csc[++cs[C2[++cr]]]++;
		while(cl>L[j]) csc[++cs[C2[--cl]]]++;
		while(cr>R[j]) csc[cs[C2[cr--]]--]--;
		res[ID[i]]=csc[K[ID[i]]];
	}
	FOR(i,M) _P("%d\n",res[i]);
}

２：頂点間のカウントマージ

当然ながら、各頂点におけるサブツリーの頂点の色の数は、その頂点の子における数の総和にその頂点自身の色の分を加えたものである。
普通に処理すると総和を取る処理がO(N)かかるが、各頂点における数をカウントする際、子のうち最大のデータ量を持つものを流用し、そこによりデータ量の小さい子の頂点のサブツリーのカウントを合わせていけばよい。
最大データを再利用することで、各頂点における総和を取る処理を軽量化できる。

そんなことを考えていたら、ちょうど下記エントリが話題になった。
"データ構造をマージする一般的なテク" とは？ - (iwi) ｛ 反省します - TopCoder部

int N,M;
int ti=0;
int C[100010],num[100010];
vector<int> E[100010];
vector<pair<int,int> > Q[100010];
int L[100010],R[100010],C2[100010];
int res[100010];
map<int,int> NM[100010];
vector<int> NV[100010];

void dfs(int cur,int pre) {
	int i,j;
	int mi=-1,ma=0;
	
	L[cur]=ti++;
	FOR(i,E[cur].size()) {
		int tar=E[cur][i];
		if(pre!=tar) {
			dfs(tar,cur);
			if(NM[tar].size()>ma) ma=NM[tar].size(),mi=tar;
		}
	}
	R[cur]=ti-1;
	num[cur]=ti-L[cur];
	
	NM[cur][C[cur]]=1;
	NV[cur].push_back(1);
	
	if(mi!=-1) {
		swap(NM[cur],NM[mi]);
		swap(NV[cur],NV[mi]);
		FOR(i,E[cur].size()) {
			int tar=E[cur][i];
			if(pre==tar) continue;
			ITR(it,NM[tar]) {
				for(j=NM[cur][it->first]+1;j<=NM[cur][it->first]+it->second;j++) {
					if(j>NV[cur].size()) NV[cur].push_back(1);
					else NV[cur][j-1]++;
				}
				NM[cur][it->first] += it->second;
			}
		}
	}
	
	FOR(i,Q[cur].size()) {
		if(NV[cur].size() >= Q[cur][i].first)
			res[Q[cur][i].second] = NV[cur][Q[cur][i].first-1];
	}
	
}

void solve() {
	int f,i,j,k,l,x,y;
	
	cin>>N>>M;
	FOR(i,N) cin>>C[i];
	FOR(i,N-1) {
		cin>>x>>y;
		E[x-1].push_back(y-1);
		E[y-1].push_back(x-1);
	}
	FOR(i,M) {
		cin>>x>>y;
		Q[x-1].push_back(make_pair(y,i));
	}
	FOR(i,N) sort(Q[i].begin(),Q[i].end());
	dfs(0,-1);
	FOR(i,M) _P("%d\n",res[i]);
}

まとめ

どちらも興味深い解法だった。
平方分割について、配列範囲に対するクエリを処理する場合、範囲の左側の移動量を減らし、右側は一方向の移動を複数回繰り返す、というやりかたは初めてだった。

後者は自分で思いついたアイデアだったけど、ちょうどそれを補強する記事が見られてよかった。
どちらも頭に入れておこう。