Codeforces #294 Div2 E. A and B and Lecture Rooms

これはなかなか面白い問題だった。
http://codeforces.com/contest/519/problem/E

問題

N頂点の木を成すグラフが与えられる。
各辺の長さは1とする。

ここでM個のクエリが与えられる。
各クエリでは2頂点A[i],B[i]が与えられるので、両頂点からの距離が等しい点の数を答えよ。

解法

基本的な考え方は、2頂点の中間の点を求め、その点にからA[i]とB[i]に至る子のサブツリーを除いた頂点数を求めることである。
そのような中間点およびサブツリーを求めるため、LCAを用いる。

まずはDFSで各頂点の親子関係と深さ、さらにサブツリーの頂点数を求める。
次にLCAを求められるようダブリングをしていく。
以下頂点xの根からの距離をD(x)、xのサブツリーの頂点数をnum(x)、xのy個親の頂点をP(x,y)とする。

以下各クエリを順次処理する。([i]は省略)

• A==Bの時は全頂点までの距離が等しいので例外として除く。
• それ以外の場合、まずL=LCA(A,B)を求める。
  • 2頂点の距離d = D(A) + D(B) - 2*D(L)が奇数の場合、条件を満たす頂点は無い。
  • それ以外の場合、AとBの中間点Mを考える。
    • D(A)==D(B)なら、M=Lである。
      • AとBを含むMの子は、P(A,D(A)-D(L)-1)、P(B,D(B)-D(L)-1)である。
      • よって解はN-num(P(A,D(A)-D(L)-1))-num(P(A,D(A)-D(L)-1))
    • D(A)＜D(B)なら、M=P(B,D(B)-d/2)である。
      • Mの親側はAに近いので条件を満たさない。
      • よって解はnum(P(B,D(B)-d/2))-num(P(B,D(B)-d/2-1))
    • D(A)＞D(B)ならAとBをswapして同じように処理すればよい。

int N,M;
vector<int> E[101000];
int num[101000];
int id;
int P[21][100005],D[100005];

int dfs(int cur,int pre) {
	int un=-1,i;
	num[cur]=1;
	FORR(tar,E[cur]) if(tar!=pre) D[tar]=D[cur]+1, P[0][tar]=cur, num[cur] += dfs(tar,cur);
	return num[cur];
}
int getpar(int cur,int up) {
	int i;
	FOR(i,20) if(up&(1<<i)) cur=P[i][cur];
	return cur;
}

int lca(int a,int b) {
	int ret=0,i,aa=a,bb=b;
	if(D[aa]>D[bb]) swap(aa,bb);
	for(i=19;i>=0;i--) if(D[bb]-D[aa]>=1<<i) bb=P[i][bb];
	for(i=19;i>=0;i--) if(P[i][aa]!=P[i][bb]) aa=P[i][aa], bb=P[i][bb];
	return (aa==bb)?aa:P[0][aa];               // vertex
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	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,19) FOR(x,N) P[i+1][x]=P[i][P[i][x]];

	cin>>M;
	while(M--) {
		cin>>x>>y;
		x--,y--;
		if(x==y) {
			cout<<N<<endl;
			continue;
		}
		if(D[x]>D[y]) swap(x,y);
		int lc=lca(x,y);
		int dist=D[x]-D[lc]+D[y]-D[lc];
		if(dist%2) {
			cout<<0<<endl;
			continue;
		}
		int ret;
		if(D[x]==D[y]) ret = N-num[getpar(x,D[x]-D[lc]-1)]-num[getpar(y,D[y]-D[lc]-1)];
		else ret = num[getpar(y,dist/2)]-num[getpar(y,dist/2-1)];
		cout<<ret<<endl;
	}
	
}

まとめ

本番はA,Bを含むサブツリーをEuler Tour+lower_boundでかなり面倒な求め方をしてしまった。
ある頂点xがサブツリー内の頂点yを含むとき、どの子頂点の先にあるかを求めるには、P(y,D(y)-D(x)-1)で求めればよいっていうのは勉強になった。