淀粉质 WA on 1 3 4 6 求调

RT

#include<bits/stdc++.h>
using namespace std;
#define N 11451
int vis[N];
bool tong[N*N];
int k;
int n,q;
vector<pair<int,int> > e[N];
int sz[N],mxs[N],rt;
vector<int> col[N];
void dfs1(int u,int fa,int n)
{
    sz[u]=1;mxs[u]=0;
    for(auto v:e[u]) if(vis[v.first]==0&&v.first!=fa) dfs1(v.first,u,n),sz[u]+=sz[v.first],mxs[u]=max(sz[v.first],mxs[u]);
    mxs[u]=max(mxs[u],n-sz[u]);
    if(mxs[u]<=n/2) rt=u;
}
void dfs2(int u,int d,int co,int fa)
{
    col[co].push_back(d);
    for(auto v:e[u]) if(v.first!=fa&&vis[v.first]==0) dfs2(v.first,d+v.second,co,u);
}
bool solve(int r,int n)
{
    if(n==1) return 0;
    dfs1(r,0,n);
    int co=0;
    for(auto v:e[rt]) if(!vis[v.first]) dfs2(v.first,v.second,++co,rt);
    tong[0]=1;
    bool ans=0;
    for(int i=1;i<=co;i++)
    {
        for(auto x:col[i]) if(x<=k&&tong[k-x]) ans=1;
        for(auto x:col[i]) tong[x]=1;
    }
    for(int i=1;i<=co;i++)
    {
        for(auto x:col[i]) tong[x]=0;
        col[i].clear();
    }
    if(ans) return ans;
    vis[rt]=1;
    for(auto v:e[rt]) if(!vis[v.first])
    {
        if(sz[v.first]<sz[rt]) ans|=solve(v.first,sz[v.first]);
        else ans|=solve(v.first,n-sz[rt]);
        if(ans) return ans;
    }
    return 0;
}
int main()
{
    cin>>n>>q;
    int u,v,w;
    for(int i=1;i<n;i++) cin>>u>>v>>w,e[u].push_back({v,w}),e[v].push_back({u,w});
    while(q--)
    {
        cin>>k;
        cout<<(solve(1,n)?"AYE":"NAY")<<endl;
        memset(vis,0,sizeof(vis));
    }
    return 0;
}