[题解]P6374 「StOI-1」树上询问

P6374 「StOI-1」树上询问

题意简述

给定一个\(N\)个节点的树，接下来有\(q\)次询问。每次询问给定\(a,b,c\)，请问存在多少个节点\(i\)，使得这棵树在以\(i\)为根的情况下，\(a\)和\(b\)的LCA是\(c\)。

解题思路 & Code

首先通过分析样例，我们发现：\(a,b\)的LCA一定在它们之间的简单路径上，所以如果\(c\)不在\(a,b\)之间的简单路径上，则输出\(0\)。

进一步分析\(c\)在\(a,b\)间简单路径上的情况，我们可以归纳出三种情况：

• 如果\(c=lca(a,b)\)，那么答案是\(n-siz[jump(a,c)]-siz[jump(b,c)]\)，其中\(siz[u]\)表示以\(u\)为根的子数大小，\(jump(u,v)\)表示\(u\)点向上跳到\(v\)的下一层所到达的节点。
• 如果\(c\)在\(a\)到\(lca(a,b)\)的路径上，那么答案是\(siz[c]-siz[jump(a,c)]\)。
• 如果\(c\)在\(b\)到\(lca(a,b)\)的路径上，那么答案是\(siz[c]-siz[jump(b,c)]\)。

接下来我们思考代码实现。怎样知道\(c\)是哪一种情况呢？

• 第一种情况：\(c=lca(a,b)\)。
• 第二种情况：\(lca(a,c)=c\)且\(lca(a,b)=lca(b,c)\)。
• 第三种情况：\(lca(b,c)=c\)且\(lca(a,b)=lca(a,c)\)。

代码实现中，\(jump()\)函数我们可以通过倍增的思想在\(O(\log N)\)的时间复杂度内完成。

LCA这里用倍增实现，预处理\(O(n\log N)\)，每次查询\(O(\log N)\)。

总时间复杂度\(O((N+q)\log N)\)。

点击查看代码

#include<bits/stdc++.h>
#define N 500010
using namespace std;
int n,q,dep[N],fa[N][20],siz[N];
vector<int> G[N];
void dfs(int u,int father){
	siz[u]=1;
	dep[u]=dep[father]+1;
	fa[u][0]=father;
	for(int i=1;i<20;i++)
		fa[u][i]=fa[fa[u][i-1]][i-1];
	for(int i:G[u])
		if(i!=father) dfs(i,u),siz[u]+=siz[i];
}
int lca(int u,int v){
	if(dep[u]<dep[v]) swap(u,v);
	for(int i=19;i>=0;i--)
		if(dep[fa[u][i]]>=dep[v])
			u=fa[u][i];
	if(u==v) return v;
	for(int i=19;i>=0;i--)
		if(fa[u][i]!=fa[v][i])
			u=fa[u][i],v=fa[v][i];
	return fa[u][0];
}
int jump(int a,int b){
	if(dep[a]==dep[b]) return 0;
	//计算a跳到b下一层的位置
	for(int i=19;i>=0;i--)
		if(dep[fa[a][i]]>dep[b])
			a=fa[a][i];
	return a;
}
int main(){
	cin>>n>>q;
	for(int i=1;i<n;i++){
		int u,v;
		cin>>u>>v;
		G[u].emplace_back(v);
		G[v].emplace_back(u);
	}
	dfs(1,0);
	while(q--){
		int a,b,c;
		cin>>a>>b>>c;
		int ab=lca(a,b),ac=lca(a,c),bc=lca(b,c);
		if(ab==c){
			cout<<n-siz[jump(a,c)]-siz[jump(b,c)]<<"\n";
		}else if(ac==c&&ab==bc){
			cout<<siz[c]-siz[jump(a,c)]<<"\n";
		}else if(bc==c&&ab==ac){
			cout<<siz[c]-siz[jump(b,c)]<<"\n";
		}else{
			cout<<"0\n";
		}
	}
	return 0;
}

附：对于上面归纳出的\(3\)种情况，给出另一种求法：

• 第一种情况：\(c=lca(a,b)\)。
• 第二种情况：\(lca(a,c)=c\)且\(dep[c]\ge dep[lca(a,b)]\)。
• 第三种情况：\(lca(b,c)=c\)且\(dep[c]\ge dep[lca(a,b)]\)。