洛谷P3806 【模板】点分治1

题目背景

感谢hzwer的点分治互测。

题目描述

给定一棵有n个点的树

询问树上距离为k的点对是否存在。

输入输出格式

输入格式：

n,m 接下来n-1条边a,b,c描述a到b有一条长度为c的路径

接下来m行每行询问一个K

输出格式：

对于每个K每行输出一个答案，存在输出“AYE”,否则输出”NAY”(不包含引号)

输入输出样例

输入样例#1： 

2 1
1 2 2
2

输出样例#1： 

AYE

说明

对于30%的数据n<=100

对于60%的数据n<=1000,m<=50

对于100%的数据n<=10000,m<=100,c<=1000,K<=10000000

题解：模板题

对于求距离为k的对数，开一个大数组储存距离,最后判断ans[k]是否存在，即可

	for(int i=1;i<=tot;i++){
		for(int j=1;j<=tot;j++)
			if(i != j)	ans[num[i]+num[j]] += flag;
	}

#include <bits/stdc++.h>
using namespace std;
int const N = 20000 + 10;
int const M = 10000000 + 10;
int const inf = 0x7f7f7f7f;
int n,m,k,tot,root,Val;
int ans[M],mav[N],size[N],num[N];
int val[2*N],first[N],ne[2*N],to[2*N];
bool vis[N];
void add(int x,int y,int dis){
	ne[++tot] = first[x];
	to[tot] = y;
	val[tot] = dis;
	first[x] = tot;
}
void Init(){
	tot = 0;
	memset(first,0,sizeof(first));
	memset(vis,false,sizeof(vis));
	for(int i=1;i<=n-1;i++){
		int x,y,dis;
		scanf("%d%d%d",&x,&y,&dis);
		add(x,y,dis);	add(y,x,dis);
	}
}
void Get_Size(int x,int fa){
	size[x] = 1,mav[x] = 0;
	for(int i=first[x];i;i=ne[i]){
		int u = to[i];
		if(u == fa || vis[u])	continue;  
		Get_Size(u,x);
		size[x] += size[u];
		mav[x] = max(mav[x],size[u]);
	}
}
void Get_Root(int sz,int x,int fa){
	mav[x] = max(size[sz] - size[x],mav[x]);
	if(mav[x] < Val){
		Val = mav[x];
		root = x;
	}
	for(int i=first[x];i;i=ne[i]){
		int u = to[i];
		if(u == fa || vis[u])	continue;
		Get_Root(sz,u,x);
	}
}
void Get_Dis(int x,int fa,int dis){
	num[++tot] = dis;
	for(int i=first[x];i;i=ne[i]){
		if(to[i] == fa || vis[to[i]])	continue;
		Get_Dis(to[i],x,dis+val[i]);
	}
}
void cacl(int x,int dis,int flag){
	int now = 0;tot = 0;
	Get_Dis(x,0,dis);
	for(int i=1;i<=tot;i++){
		for(int j=1;j<=tot;j++)
			if(i != j)	ans[num[i]+num[j]] += flag;
	}
}
void DFS(int x){
	Val = inf;
	Get_Size(x,0);
	Get_Root(x,x,0);
	cacl(root,0,1);
	vis[root] = true;	
	for(int i=first[root];i;i=ne[i]){
		if(!vis[to[i]]){
			cacl(to[i],val[i],-1);
			DFS(to[i]);
		}
	}
}
int main(){
	cin>>n>>m;
	Init();
	DFS(1);
	for(int i=1;i<=m;i++){
		int tmp;
		scanf("%d",&tmp);
		if(ans[tmp])	printf("AYE\n");
		else	printf("NAY\n");
	}
}