点分治模板。把询问离线下来回答,记一下每种长度的出现次数就好,没什么需要特别注意的。
#include<bits/stdc++.h>
using namespace std;
const int maxn=1e5+10;
int Head[maxn],Next[maxn<<1],V[maxn<<1],W[maxn<<1],cnt=0;
int siz[maxn],mx[maxn],vis[maxn],dis[maxn],ask[maxn],ret[maxn],SZ,root=0,tot=0;
int n,m,a,b,c,k;
inline int read(){
int x=0;char ch=getchar();
while(!isdigit(ch)) ch=getchar();
while(isdigit(ch)) x=(x<<3)+(x<<1)+ch-'0',ch=getchar();
return x;
}
inline void add(int u,int v,int w){
++cnt;
Next[cnt]=Head[u];
V[cnt]=v,W[cnt]=w;
Head[u]=cnt;
}
void get_root(int u,int f){
mx[u]=0,siz[u]=1;
for(int i=Head[u];i;i=Next[i]) if((!vis[V[i]])&&(V[i]!=f)){
get_root(V[i],u);
mx[u]=max(mx[u],siz[V[i]]);
siz[u]+=siz[V[i]];
}
mx[u]=max(mx[u],SZ-siz[u]);
if(mx[root]>mx[u]) root=u;
}
void get_dis(int u,int f,int D){
for(int i=Head[u];i;i=Next[i]) if((!vis[V[i]])&&(V[i]!=f)){
dis[++tot]=D+W[i];
get_dis(V[i],u,dis[tot]);
}
}
int get_ans(int u,int D,int k,int l=1,int ret=0){
dis[tot=1]=D,get_dis(u,0,D);
sort(dis+1,dis+tot+1);
while(l<tot && dis[l]+dis[tot]<k) ++l;
while(l<tot && dis[l]<=(k-dis[l])){
int p1=lower_bound(dis+l+1,dis+tot+1,k-dis[l])-dis;
int p2=upper_bound(dis+l+1,dis+tot+1,k-dis[l])-dis;
ret+=p2-p1,++l;
}
return ret;
}
void dfs(int u){
vis[u]=1;
for(int i=1;i<=m;++i) ret[i]+=get_ans(u,0,ask[i]);
for(int i=Head[u];i;i=Next[i]) if(!vis[V[i]]){
for(int j=1;j<=m;++j) ret[j]-=get_ans(V[i],W[i],ask[j]);
root=0,SZ=siz[V[i]],get_root(V[i],u),dfs(root);
}
}
int main(){
SZ=mx[0]=n=read(),m=read();
for(int i=1;i<n;++i) a=read(),b=read(),c=read(),add(a,b,c),add(b,a,c);
for(int i=1;i<=m;++i) ask[i]=read();
get_root(1,0),dfs(root);
for(int i=1;i<=m;++i)
if(ret[i]>0) puts("AYE");
else puts("NAY");
}