树上路径问题点分治

点分治
点分治主要解决树上的路径问题
用到了树的重心

#include<iostream>
using namespace std;
const int maxn = 2e7+7;
const int N = 2e4+7;
int ext[maxn],rem[N];
int dis[N],query[N],root,vis[N],used[N],max_size;
int INF = 0x3f3f3f3f,num1,num2,ans[maxn];
int n,m;
struct Edge{
	int to,next,w;
}edge[maxn];
int head[maxn],cnt=0,sz[maxn];

void add(int u,int v,int w){
	edge[++cnt].to = v;
	edge[cnt].w = w;
	edge[cnt].next = head[u];
	head[u] = cnt;
}

void get_rt(int u,int fa,int n){
	sz[u] = 1;
	int tmp = 0;
	for(int i=head[u];i;i=edge[i].next){
		int v = edge[i].to;
		if(v == fa || vis[v])  continue;
		get_rt(v,u,n);
		sz[u] += sz[v];
		tmp = max(tmp,sz[v]);
	}
	tmp = max(tmp,n-sz[u]);
	if(tmp < max_size){
		max_size = tmp;
		root = u;
	}
}

void get_v_rt_dis(int u,int fa){
	rem[++num1] = dis[u];
	for(int i=head[u];i;i=edge[i].next){
		int v = edge[i].to;
		int w = edge[i].w;
		if(v == fa || vis[v])  continue;
		dis[v] = dis[u] + w;
		get_v_rt_dis(v,u);
	}
}

void cal_exist(int u){
	num2 = 0;
	for(int i=head[u];i;i=edge[i].next){
		int v = edge[i].to;
		if(vis[v])   continue;
		num1 = 0,dis[v] = edge[i].w;
		get_v_rt_dis(v,u);
		for(int j=1;j<=num1;j++){
			for(int k=1;k<=m;k++){
				int query_len = query[k] - rem[j];
				if(query_len >= 0 && query_len <= 1e7 && ext[query_len]){
					ans[k] = 1;
				}
			}
		}
		for(int j=1;j<=num1;j++){
			if(rem[j] <= 1e7){
				used[++num2] = rem[j];
				ext[rem[j]] = 1;
			}
		}
	}
	for(int i=1;i<=num2;i++){
		ext[used[i]] = 0;	
	}
}

void divide(int u){
	vis[u] = 1;
	ext[0] = 1;
	cal_exist(u);
	for(int i=head[u];i;i=edge[i].next){
		int v = edge[i].to;
		if(vis[v])   continue;
		max_size = INF;
		get_rt(v,u,sz[v]);
		divide(root);
	}
}

int main(){
	scanf("%d %d",&n,&m);
	for(int i=1;i<=n-1;i++){
		int u,v,w;
		scanf("%d %d %d",&u,&v,&w);
		add(u,v,w);
		add(v,u,w);
	}
	for(int i=1;i<=m;i++){
		scanf("%d",&query[i]);
	}
	max_size = INF;
	get_rt(1,0,n);
	divide(root);
	for(int i=1;i<=m;i++){
		if(ans[i]){
			puts("AYE");
		}
		else{
			puts("NAY");
		}
	}
}