[Luogu3806]点分治

Posted 2020-11-09 void_f

询问树上是否存在距离为k[i]的点对

直接点分治把所有距离预处理出来，然后O(1)回答即可

Code

#include <cstdio>
#include <algorithm>
#define N 10010
using namespace std;

const int mx=N*1000;
struct info{int to,nex,w;}e[N<<1];
int n,m,tot,head[N],Ans[mx],sz[N],rt,d[N],sum,f[N];
bool vis[N];

void Link(int u,int v,int w){
	e[++tot].to=v,e[tot].w=w,e[tot].nex=head[u];head[u]=tot;
}

inline int read(){
    int x=0,f=1;char ch=getchar();
    while(ch<‘0‘||ch>‘9‘){if(ch==‘-‘)f=-1;ch=getchar();}
    while(ch>=‘0‘&&ch<=‘9‘){x=x*10+ch-‘0‘;ch=getchar();}
    return x*f;
}

void getrt(int u,int fa){
	sz[u]=1,f[u]=0;
	for(int i=head[u];i;i=e[i].nex){
		int v=e[i].to;
		if(v==fa||vis[v]) continue;
		getrt(v,u);
		sz[u]+=sz[v];
		f[u]=max(f[u],sz[v]);
	}
	f[u]=max(f[u],sum-sz[u]);
	if(f[rt]>f[u]) rt=u;
} 

void getdep(int u,int fa,int dep){
	d[++d[0]]=dep;
	for(int i=head[u];i;i=e[i].nex){
		int v=e[i].to;
		if(v==fa||vis[v]) continue;
		getdep(v,u,dep+e[i].w);
	}
}

void calc(int u,int f,int pre){
	d[0]=0;
	getdep(u,0,0);
	for(int i=1;i<=d[0];++i)
		for(int j=i+1;j<=d[0];++j)
			if(f&&d[i]+d[j]<=mx) ++Ans[d[i]+d[j]];
			else if(d[i]+d[j]+pre<=mx) --Ans[d[i]+d[j]+pre];
}

void solve(int u){
	calc(u,1,0);
	vis[u]=1;
	for(int i=head[u];i;i=e[i].nex){
		int v=e[i].to;
		if(vis[v]) continue;
		calc(v,0,e[i].w*2);
		sum=sz[v];
		getrt(v,rt=0);
		solve(rt);
	}
}

int main(){
	n=read(),m=read();
	for(int i=1;i<n;++i){
		int u=read(),v=read(),w=read();
		Link(u,v,w),Link(v,u,w);
	}
	sum=n,f[0]=1e9;
	getrt(1,0);
	solve(rt);
	while(m--){
		int x=read();
		puts(Ans[x]?"AYE":"NAY");
	}
	return 0;
}