こちらはちょっと時間食ったけど方針としてはまぁすんなり。
https://csacademy.com/contest/round-63/task/root-lca-queries/
問題
木を成すグラフが与えられる。
この木に対し、以下のクエリに答えよ。
3頂点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の良い練習問題。