T了40分，感觉没毛病啊qwq

那位大佬帮我优化一下呗qwq

    #include <bits/stdc++.h>
    #define nmax 1000000
    using namespace std;
    struct edge
    {
        int to;
        int lon;
        edge(int to=-1,int lon=0):to(to),lon(lon) {}
    };
    vector<edge> G[nmax];
    pair<int,int> v[nmax];
    bool vis[nmax];
    map<int,bool> ans;
    int root,siz[nmax],rootmax,p;
    void get(int x,int dis,int color,int from)
    {
//      printf("%d %d %d %d\n",x,dis,color,from);
        v[p++]=make_pair(dis,color);
        for(int i=0;i<G[x].size();i++)
            if(!vis[G[x][i].to] && G[x][i].to!=from)
                get(G[x][i].to , dis+G[x][i].lon , color , x);
    }
    void getroot(int x,int from,int n)
    {
        int max_=0;
        siz[x]=1;
        for(int i=0;i<G[x].size();i++)
            if(G[x][i].to != from && !vis[G[x][i].to])
            {
                getroot(G[x][i].to,x,n);
                max_=max(siz[G[x][i].to] , max_);
                siz[x]+=siz[G[x][i].to];
            }
        max_=max(n-siz[x] , max_);
        if(max_ < rootmax) root=x , rootmax=max_;
    }
    void work(int x,int n)
    {
//      printf("%d",x);
        int color=0;
        rootmax=100000000;
        getroot(x,-1,n);
        x=root;
        p=0;
        v[p++]=make_pair(0,0);
//      printf("->%d\n",x);
        for(int i=0;i<G[x].size();i++)
            if(!vis[G[x][i].to])
                get(G[x][i].to , G[x][i].lon , ++color , x);
//      sort(v.begin() , v.end());
        for(int i=0;i<p;i++)
            for(int j=i+1;j<p;j++)
                if(v[i].second != v[j].second)
                    ans[v[i].first+v[j].first] = true;
        vis[x]=1;
        for(int i=0;i<G[x].size();i++)
            if(!vis[G[x][i].to])
                work(G[x][i].to , (siz[G[x][i].to]>siz[x]) ? n-siz[x] : siz[G[x][i].to]);
    }   
    int main()
    {
        int n,m;
        scanf("%d %d",&n,&m);
        for(int i=1;i<n;i++)
        {
            int a,b,c;
            scanf("%d %d %d",&a,&b,&c);
            G[a].push_back(edge(b,c));
            G[b].push_back(edge(a,c));
        }
        work(1,n);
        while(m--)
        {
            int need;
            scanf("%d",&need);
            if(ans[need]) printf("AYE\n");
            else printf("NAY\n");
        }
    }