P3806 【模板】点分治1

#include<bits/stdc++.h>
using namespace std;
#define rep(i,a,b) for(int i=(a);i<=(b);i++)
#define repp(i,a,b) for(int i=(a);i>=(b);--i)
#define ll long long
#define see(x) (cerr<<(#x)<<‘=‘<<(x)<<endl)
#define inf 0x3f3f3f3f
#define CLR(A,v)  memset(A,v,sizeof A)
//////////////////////////////////
const int N=1e5+10;
int n,m,ans[N],a[N],num,root,SIZE,siz[N],pos,head[N],maxx,f[N],vis[N];
struct Edge{int to,nex,w;}edge[N];
void add(int a,int b,int c)
{
    edge[++pos]=(Edge){b,head[a],c};
    head[a]=pos;
}
void getroot(int x,int fa)
{
    f[x]=0;siz[x]=1;
    for(int i=head[x];i;i=edge[i].nex)
    {
        int v=edge[i].to;if(vis[v]||v==fa)continue;
        getroot(v,x);
        f[x]=max(f[x],siz[v]);
        siz[x]+=siz[v];
    }
    f[x]=max(f[x],SIZE-siz[x]);
    if(f[x]<f[root])root=x;
}
struct node
{
    int id,v;
}dis[N];
void dfs(int x,int fa,int id,int d)
{
    dis[++num]=(node){id,d};
    for(int i=head[x];i;i=edge[i].nex)
    {
        int v=edge[i].to;
        if(vis[v]||v==fa)continue;
        dfs(v,x,id,d+edge[i].w);
    }
}
int find1(int x)
{
    int ans=0,L=1,R=num;
    while(L<=R)
    {
        int mid=(L+R)>>1;
        if(dis[mid].v>=x)ans=mid,R=mid-1;
        else L=mid+1;
    }
    return ans;
}
void work(int x)
{
    num=0;
    for(int i=head[x];i;i=edge[i].nex)
    {
        int v=edge[i].to;
        if(vis[v])continue;
        dfs(v,x,v,edge[i].w);
    }
    dis[++num]=(node){0,0};
    sort(dis+1,dis+1+num,[](node a,node b){return a.v<b.v;});
    rep(i,1,m)
    {
        if(ans[i])continue;
        int L=1;
        while(L<num&&dis[L].v+dis[num].v<a[i])L++;
        while(L<num&&!ans[i])
        {
            if(2*dis[L].v>a[i])break;
            int pos=find1(a[i]-dis[L].v);
            while(L<=num&&dis[pos].id==dis[L].id&&dis[pos].v+dis[L].v==a[i])pos++;
            if(dis[pos].v+dis[L].v==a[i])ans[i]=1;
            L++;
        }
    }
}
void solve(int x)
{
    vis[x]=1;
    work(x);
    for(int i=head[x];i;i=edge[i].nex)
    {
        int v=edge[i].to;
        if(vis[v])continue;
        root=0;
        siz[root]=SIZE=siz[v];
        getroot(v,0);
        solve(v);
    }
}
int main()
{
    cin>>n>>m;
    rep(i,1,n-1)
    {
        int a,b,c;cin>>a>>b>>c;
        add(a,b,c);add(b,a,c);
    }
    rep(i,1,m)cin>>a[i];

    f[root]=SIZE=n;
    getroot(1,0);
    solve(root);
    rep(i,1,m)
    if(ans[i])printf("AYE\n");
    else printf("NAY\n");
    return 0;
}