CSAcademy Round #63 : E. Root LCA Queries

こちらはちょっと時間食ったけど方針としてはまぁすんなり。
https://csacademy.com/contest/round-63/task/root-lca-queries/

問題

木を成すグラフが与えられる。
この木に対し、以下のクエリに答えよ。

３頂点A,B,Cが指定される。
この木の根となる頂点をDとするとき、AとBのLCAがCとなるようなDはいくつあるか。

なお、A,B,Cは互いに異なる頂点である。

解法

まずは根頂点を仮に設定し、EulerTourとLCAを求めるためのダブリングを済ませておこう。

まず条件として、CはA-Bの最短経路上になければならない。
A-Bの最短経路上以外にLCAがあることはありえない。
これは以下のいずれかを満たすかをチェックすればよい。どちらも満たさないなら解は0。

• CがLCA(A,B)と一致する
• CがLCA(A,B)よりも深いところにあり、かつAかBの祖先である

後者の祖先判定は最初に頂点の深さを求めEulerTourをしておけば容易に判定できる。

さてそれぞれの場合にDの候補を数え上げよう。

• CがLCA(A,B)と一致する場合：
  • Cより1つA寄りの点をA'、B寄りの点をB'とする。
  • A'のSubTreeとB'のSubTreeにDがある場合、条件を満たさない。よってNから両SubTreeの頂点数を引けばよい。
• CがLCA(A,B)よりも深いところにあり、かつAかBの祖先である場合：
  • 以下CがAの祖先である場合を考える。
  • DはCのSubTree内でかつA'のSubTreeに無ければいいので、前者から後者を引けばよい。

int N,Q;
vector<int> E[101010];
int id;
int L[101010],R[101010],V[101010],rev[101010];
int P[21][200005],D[200005];
int A,B,C;

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 dfs(int cur) {
	L[cur]=id++;
	V[cur]=1;
	rev[L[cur]]=cur;
	ITR(it,E[cur]) if(*it!=P[0][cur]) D[*it]=D[cur]+1, P[0][*it]=cur, dfs(*it), V[cur]+=V[*it];
	R[cur]=id;
}

int getpar(int cur,int up) {
	int i;
	FOR(i,20) if(up&(1<<i)) cur=P[i][cur];
	return cur;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>Q;
	FOR(i,N-1) {
		cin>>x>>y;
		E[x-1].push_back(y-1);
		E[y-1].push_back(x-1);
	}
	
	dfs(0);
	FOR(i,19) FOR(x,N) P[i+1][x]=P[i][P[i][x]];
	
	while(Q--) {
		cin>>A>>B>>C;
		A--,B--,C--;
		int lc=lca(A,B);
		if(C==lc) {
			int a2=getpar(A,D[A]-D[C]-1);
			int b2=getpar(B,D[B]-D[C]-1);
			cout<<N-V[a2]-V[b2]<<endl;
		}
		else if(D[C]>=D[lc] && L[C]<=L[A] && L[A]<R[C]) {
			cout<<V[C]-V[getpar(A,D[A]-D[C]-1)]<<endl;
		}
		else if(D[C]>=D[lc] && L[C]<=L[B] && L[B]<R[C]) {
			cout<<V[C]-V[getpar(B,D[B]-D[C]-1)]<<endl;
		}
		else {
			cout<<0<<endl;
		}
	}
}

まとめ

EulerTourとLCAの良い練習問題。