萌新刚学点分治，求助

评测记录https://www.luogu.com.cn/record/31481927

#include <iostream>
#include <cstdio>
using namespace std;

const int MAXN=10005,MAXF=20000005,INF=0x3f3f3f3f;
struct node{
    int v,c;
    node*next;
}*h[MAXN],pool[MAXN<<1];
int dis[MAXN],dis2[MAXN],siz[MAXN],q[MAXN],maxa[MAXN],qst[MAXN];
int n,m,rt,sum,tot=0,cnt=0,ctt=0;
bool vis[MAXN],ans[MAXN],c[MAXF];

void addedge(int u,int v,int c){
    node*p=&pool[++tot];
    p->v=v;p->c=c;p->next=h[u];h[u]=p;
}

void getrt(int u,int fa){
    siz[u]=1;
    maxa[u]=-1;
    for(node*p=h[u];p;p=p->next){
        if(p->v==fa||vis[p->v]){
            continue;
        }
        getrt(p->v,u);
        siz[u]+=siz[p->v];
        maxa[u]=max(maxa[u],siz[p->v]);
    }
    maxa[u]=max(maxa[u],sum-siz[u]);
    if(maxa[u]<maxa[rt]){
        rt=u;
    }
}

void getdis(int u,int fa){
    cnt++;
    dis2[cnt]=dis[u];
    for(node*p=h[u];p;p=p->next){
        if(p->v==fa||dis[p->v]){
            continue;
        }
        dis[p->v]=dis[u]+p->c;
        getdis(p->v,u);
    }
}

void query(int u){
    vis[u]=1;
    c[0]=1;
    ctt=0;
    for(node*p=h[u];p;p=p->next){
        if(vis[p->v]){
            continue;
        }
        cnt=0;
        dis[p->v]=p->c;
        getdis(p->v,u);
        for(int i=1;i<=cnt;i++){
            for(int j=1;j<=m;j++){
                if(qst[j]>=dis2[i]){
                    if(c[qst[j]-dis2[i]]){
                        ans[j]=1;
                    }
                }
            }
        }
        for(int i=1;i<=cnt;i++){
            ctt++;
            q[ctt]=dis2[i];
            c[dis2[i]]=1;
        }
    }
    for(int i=1;i<=ctt;i++){
        c[q[i]]=0;
    }
    for(node*p=h[u];p;p=p->next){
        if(vis[p->v]){
            continue;
        }
        sum=siz[p->v];
        rt=0;
        maxa[0]=INF;
        getrt(p->v,0);
        query(rt);
    }
}

int main(){
    scanf("%d%d",&n,&m);
    for(int i=1;i<n;i++){
        int u,v,c;
        scanf("%d%d%d",&u,&v,&c);
        addedge(u,v,c);
        addedge(v,u,c);
    }
    for(int i=1;i<=m;i++){
        scanf("%d",&qst[i]);
    }
    rt=0;
    sum=n;
    maxa[rt]=n;
    getrt(1,0);
    query(rt);
    for(int i=1;i<=m;i++){
        if(ans[i]){
            printf("AYE\n");
        }
        else{
            printf("NAY\n");
        }
    }
    return 0;
}