萌新的蜜汁被卡TLE #7 #9

被卡常了 但是这份代码是过了http://poj.org/problem?id=2114 这道题的(几乎是原题)

本地跑#7 100ms 左右

#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
const int M=1e4+10;

struct node{
    int to,next,dis;
}edge[M<<1];
int head[M],cnte;
void add(int u,int v,int w){
    edge[++cnte].to=v;
    edge[cnte].dis=w;
    edge[cnte].next=head[u];
    head[u]=cnte;
}

int n,m,ans,k;

bool vis[M];
int size[M],f[M],tot,root;
void getroot(int u,int fa){
    size[u]=1,f[u]=0;
    for(int i=head[u];~i;i=edge[i].next){
        int v=edge[i].to;
        if(v!=fa&&!vis[v]){
            getroot(v,u);
            size[u]+=size[v];
            f[u]=max(f[u],size[v]);
        }
    }
    f[u]=max(f[u],tot-size[u]);
    if(f[u]<f[root]) root=u;
}

int dep[M],d[M];
void getdep(int u,int fa){
    dep[++dep[0]]=d[u];
    for(int i=head[u];~i;i=edge[i].next){
        int v=edge[i].to;
        if(v!=fa&&!vis[v]){
            d[v]=d[u]+edge[i].dis;
            getdep(v,u);
        }
    }
}

int getans(int u,int dis){
    d[u]=dis,dep[0]=0;
    getdep(u,0);
    sort(dep+1,dep+1+dep[0]);
    int l=1,r=dep[0],cnt=0;
    while(l<r){
        if(dep[l]+dep[r]<k) l++;
        else if(dep[l]+dep[r]>k) r--;
        else{
            if(dep[l]==dep[r]){
                cnt+=(r-l+1)*(r-l)/2;
                break;
            }
            int st=l,ed=r;
            while(dep[st]==dep[l]) st++;
            while(dep[ed]==dep[r]) ed--;
            cnt+=(st-l)*(r-ed);
            l=st,r=ed;
        }
    }
    return cnt;
}

void solve(int u){
    vis[u]=true,ans+=getans(u,0);
    for(int i=head[u];~i;i=edge[i].next){
        int v=edge[i].to;
        if(!vis[v]){
            ans-=getans(v,edge[i].dis);
            root=0,tot=size[v];
            getroot(v,u),solve(root);
        }
    }
}

int main(){
    memset(head,-1,sizeof head);
    scanf("%d%d",&n,&m);
    for(int i=1,u,v,w;i<n;i++){
        scanf("%d%d%d",&u,&v,&w);
        add(u,v,w),add(v,u,w);
    }
    for(int i=1;i<=m;i++){
        for(int i=0;i<=n;i++) vis[i]=false;
        f[0]=2147483647;
        root=ans=0,tot=n;
        scanf("%d",&k);
        getroot(1,0),solve(root);
        if(ans) puts("AYE");
        else puts("NAY");
    }
    return 0;
}