60pts TLE #7 #9 求调

#include <bits/stdc++.h>
using namespace std;

const int N=10005;
const int M=20005;
const int S=10000005;

int n,m,K;
int e[M],ne[M],w[M],h[N],tot;
bool del[N];
int arr1[N],idx1,arr2[N],idx2;

void add(int a,int b,int c) {
    e[tot]=b,w[tot]=c,ne[tot]=h[a],h[a]=tot++;
}

int get_size(int u,int fa) {
    if (del[u]) return 0;
    int ans=1;
    for (int i=h[u];~i;i=ne[i]) {
        if (e[i]==fa) continue;
        ans+=get_size(e[i],u);
    }
    return ans;
}

int get_center(int u,int fa,int tot,int& cen) {
    if (del[u]) return 0;
    int sum=1,mx=0;
    for (int i=h[u];~i;i=ne[i]) {
        if (e[i]==fa) continue;
        int ret=get_center(e[i],u,tot,cen);
        sum+=ret;
        mx=max(mx,ret);
    }
    mx=max(mx,tot-sum);
    if (mx<=tot/2) {
        cen=u;
    }
    return sum;
}

void get_dis(int u,int fa,int dis,int* arr,int& idx) {
    if (del[u]) return;
    arr[++idx]=dis;
    for (int i=h[u];~i;i=ne[i]) {
        if (e[i]==fa) continue;
        get_dis(e[i],u,dis+w[i],arr,idx);
    }
}

bool cmp(const int& a,const int& b) {
    return a<b;
}

int cnt[S];

int calc(int* arr,int idx) {
    for (int i=1;i<=idx;++i) {
        if (arr[i]>K) continue;
        cnt[arr[i]]++;
    }
    int ans=0;
    for (int i=1;i<=idx;++i) {
        int t=K-arr[i];
        if (t<0) continue;
        else if (!cnt[t]) continue;
        else if (t==arr[i]) ans+=cnt[t]*(cnt[t]-1)/2;
        else ans+=cnt[arr[i]]*t;
    }
    for (int i=1;i<=idx;++i) {
        if (arr[i]>K) continue;
        cnt[arr[i]]--;
    }
    return ans;
}

int solve(int u) {
    if (del[u]) return 0;
    get_center(u,0,get_size(u,0),u);
    del[u]=true;
    int res=0;
    idx1=0;
    for (int i=h[u];~i;i=ne[i]) {
        idx2=0; get_dis(e[i],u,w[i],arr2,idx2);
        res-=calc(arr2,idx2);
        for (int j=1;j<=idx2;++j) {
            if (arr2[j]==K) res++;
            arr1[++idx1]=arr2[j];
        }
    }
    res+=calc(arr1,idx1);
    for (int i=h[u];~i;i=ne[i]) res+=solve(e[i]);
    return res;
}

int main() {
    // freopen("P3806_7.in","r",stdin);
    // freopen("P3806_7_my.out","w",sdtdout);
    // clock_t begin=clock();
    memset(h,-1,sizeof(h));
    scanf("%d %d",&n,&m);
    for (int i=1;i<=n-1;++i) {
        int a,b,c;
        scanf("%d %d %d",&a,&b,&c);
        add(a,b,c);
        add(b,a,c);
    }
    while (m--) {
        for (int i=1;i<=n;++i) del[i]=false;
        scanf("%d",&K);
        int ans=solve(1);
        if (ans) puts("AYE");
        else puts("NAY");
    }
    // clock_t end=clock();
    // printf("%d\n",end-begin);
    return 0;
}