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

难度：提高

算法：树上最近公共祖先LCA

实现：树上最近公共祖先LCA and 分类讨论

题目描述：

给定一棵 n 个点的无根树，有 q 次询问。每次询问给一个参数三元组 (x,y,v)，求有多少个 u 满足这棵树在以 u 为根的情况下 x 和 y 的 LCA 为 v。

solution:

1. 当 v=LCA(x,y) 时 答案为 n−size[root[a]]-size[root[b]]

2.当 v 在 x -> LCA(x,y) 的路径上时 答案为 size[v]-size[root[x]]

3.当 v 在 y -> LCA(x,y) 的路径上时 答案为 size[v]-size[root[y]]

else.当v不在 x -> LC A(x,y) -> y 的路径上时 答案为 0

#include<bits/stdc++.h>
using namespace std;
const int N=5e5+5;
int n,m;
int dep[N],sz[N],fa[N][25];
vector<int> g[N];
void dfs(int u,int par){
	dep[u]=dep[par]+1;
	sz[u]=1;
	for(int i=1;i<=20;i++){
		fa[u][i]=fa[fa[u][i-1]][i-1];
	}
	for(auto v:g[u]){
		if(v==par){
			continue;
		}
		fa[v][0]=u;
		dfs(v,u);
		sz[u]+=sz[v];
	}
}
int lca(int x,int y){
	if(dep[x]<dep[y]){
		swap(x,y);
	}
	for(int i=20;i>=0;i--){
		if(dep[fa[x][i]]>=dep[y]){
			x=fa[x][i];
		}
	}
	if(x==y){
		return x;
	}
	for(int i=20;i>=0;i--){
		if(fa[x][i]!=fa[y][i]){
			x=fa[x][i];
			y=fa[y][i];
		}
	}
	return fa[x][0];
}
bool check(int x,int y,int u){
	int lxy=lca(x,y);
	return (lca(x,u)==u||lca(y,u)==u)&&lca(lxy,u)==lxy;
}
int root(int x,int y){
    if(x==y){
        return 0;
    }
	if(dep[x]<dep[y]){
		swap(x,y);
	}
	for(int i=20;i>=0;i--){
		if(dep[fa[x][i]]>dep[y]){
			x=fa[x][i];
		}
	}
	return sz[x];
}
int main(){
	cin>>n>>m;
	for(int i=1;i<n;i++){
		int x,y;
		cin>>x>>y;
		g[x].push_back(y);
		g[y].push_back(x);
	}
	dfs(1,0);
	while(m--){
		int x,y,u;
		cin>>x>>y>>u;
		if(lca(x,y)==u){
			cout<<n-root(x,u)-root(y,u)<<endl; 
		}
		else if(check(x,lca(x,y),u)){
			cout<<sz[u]-root(x,u)<<endl;
		}
		else if(check(y,lca(x,y),u)){
			cout<<sz[u]-root(y,u)<<endl;
		}
		else{
			cout<<"0"<<endl;
		}
	}
	return 0;
}