P3806 【模板】点分治1

P3806 【模板】点分治1

关键点

1. 设置siz[1]=n; in[0]=1; 注意数组的更新。
2. 在对子树递归时不要访问到子树之外，对一棵子树的递归不会改变其他子树的原始信息

#include<iostream>
#include<cstdio>
#include<vector>
#include<algorithm>
#include<cmath>
#define inf 1e9
using namespace std;
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;
}
const int maxn=20005;
int n,m;
int qry[maxn];
int h[maxn],to[maxn],nxt[maxn],tot,w[maxn];
void ade(int x,int y,int z){
	to[++tot]=y; nxt[tot]=h[x]; h[x]=tot; w[tot]=z;
}

bool had[maxn],ans[maxn];
int siz[maxn],rt,cnt,mn,dis[maxn];
bool in[10000005];
int stk[maxn],top;
void getrt(int f,int u){
	siz[u]=1;
	int mx=-inf;
	for(int i=h[u];i;i=nxt[i]){
		int v=to[i];
		if(v==f||had[v]) continue;
		getrt(u,v);
		siz[u]+=siz[v];
		mx=max(mx,siz[v]);
	}
	mx=max(mx,cnt-siz[u]);
	if(mx<mn) rt=u,mn=mx;
}
void upddis(int f,int u){
	for(int i=h[u];i;i=nxt[i]){
		int v=to[i];
		if(v==f||had[v]) continue;
		dis[v]=dis[u]+w[i];
		upddis(u,v);
	}
}
void putin(int f,int u){
	stk[++top]=u;
	for(int i=1;i<=m;i++){
		if(qry[i]-dis[u]>=0&&in[qry[i]-dis[u]]) ans[i]=1;
	}
	for(int i=h[u];i;i=nxt[i]){
		int v=to[i];
		if(v==f||had[v]) continue;
		putin(u,v);
	}
}
void calc(int u){
	dis[u]=0;
	upddis(0,u);
	for(int i=h[u];i;i=nxt[i]){
		int v=to[i];
		if(had[v]) continue;
		int st=top+1;
		in[0]=1; putin(0,v);
		for(int j=st;j<=top;j++) in[dis[stk[j]]]=1;
	}
	while(top) in[dis[stk[top]]]=0,top--;
}
void solve(int u){
	mn=inf; cnt=siz[u];
	getrt(0,u);
	had[rt]=true;
	calc(rt);
	for(int i=h[rt];i;i=nxt[i]){
		int v=to[i];
		if(had[v]) continue;
		solve(v);
	}
}
int main(){
//	freopen("c.in","r",stdin);
//	freopen("c.out","w",stdout);
	n=read(); m=read();
	for(int i=1;i<n;i++){
		int x=read(),y=read(),z=read();
		ade(x,y,z); ade(y,x,z);
	}
	for(int i=1;i<=m;i++) qry[i]=read();
	siz[1]=n; solve(1);
	for(int i=1;i<=m;i++) {
		if(ans[i]) puts("AYE");
		else puts("NAY");
	}
	return 0;
}