洛谷 P3806 【模板】点分治1

#include<stdio.h>
#include<string.h>
#include<algorithm>
#define K 10000005
#define N 10005
using std::max;
int n,m;
int hd[N],nxt[N<<1],to[N<<1],w[N<<1];
int vis[N];
int rt,sz[N],siz[N];
int son[N],dfn[N],cnt;
bool mp[K];
void add(int a,int b,int c) {
    static int tot;
    tot++;
    to[tot]=b, w[tot]=c, nxt[tot]=hd[a], hd[a]=tot;
}
void getrt(int u,int f) {
    sz[u]=1;
    for(int i=hd[u];i;i=nxt[i]) {
        if(to[i]!=f&&!vis[to[i]]) {
            getrt(to[i],u);
            siz[u]=max(siz[u],sz[to[i]]);
            sz[u]+=sz[to[i]];
        }
    }
    siz[u]=max(siz[u],n-sz[u]);
    if(!rt||siz[rt]>siz[u]) rt=u;
}

void getdis(int u,int f,int g,int dep) {
    dfn[++cnt]=dep, son[cnt]=g;
    for(int i=hd[u];i;i=nxt[i]) {
        if(to[i]!=f&&!vis[to[i]]) {
            getdis(to[i],u,g,dep+w[i]);
        }
    }
}

void solve(int u) {
    rt=0;
    getrt(u,0);
    cnt=0;
    for(int i=hd[rt];i;i=nxt[i]) {
        if(!vis[to[i]]) {
            getdis(to[i],rt,to[i],w[i]);
        }
    }
    for(int i=0;i<=cnt;i++) {
        for(int j=0;j<=cnt;j++) {
            if(son[i]!=son[j]&&dfn[i]+dfn[j]<K) {
                mp[dfn[i]+dfn[j]]=1;
            }
        }
    }
    vis[rt]=1;
    for(int i=hd[rt];i;i=nxt[i]) {
        if(!vis[to[i]]) {
            solve(to[i]);
        }
    }
}

int main() {
    int a,b,c;
    scanf("%d%d",&n,&m);
    for(int i=1;i<n;i++) {
        scanf("%d%d%d",&a,&b,&c);
        add(a,b,c), add(b,a,c);
    }
    solve(1);
    while(m--) {
        scanf("%d",&a);
        puts(mp[a] ? "AYE" : "NAY");
    }
    
}