POJ 1741 Tree

http://poj.org/problem?id=1741

点分治QAQ，测了20min

#include <cstdio>
#include <cstring>
#include <algorithm>
#define INF 0x7fffffff
const int maxn = 10007;
inline int read() {
    int x=0;char c=getchar();
    while(c<'0'||c>'9') c=getchar();
    while(c<='9'&&c>='0')x=x*10+c-'0',c=getchar();
    return x;
}
int n,k;
bool vis[maxn];
int sum,ans,root,num;
int head[10005],deep[10005],dis[10005],f[10005],son[10005];
struct node{
    int v,next,w;
}edge[maxn<<1];
void add_edge(int u,int v,int w) {
    edge[++num].v=v;edge[num].w=w;edge[num].next=head[u];head[u]=num;
}
void getroot(int x,int fa) {
    son[x]=1;f[x]=0;
    for(int i=head[x];i;i=edge[i].next) {
        int v=edge[i].v;
        if(v==fa||vis[v])continue;
        getroot(v,x);
        son[x]+=son[v];
        f[x]=std::max(f[x],son[v]);
    }
    f[x]=std::max(f[x],sum-son[x]);
    if(f[x]<f[root])root=x;
}
void getdeep(int x,int fa) {
    deep[++deep[0]]=dis[x];
    for(int i=head[x];i;i=edge[i].next) {
        int v=edge[i].v;
        if(v==fa||vis[v])continue;
        dis[v]=dis[x]+edge[i].w;
        getdeep(v,x);
    }
}
int calc(int x,int now) {
    dis[x]=now;deep[0]=0;
    getdeep(x,0);
    std::sort(deep+1,deep+deep[0]+1);
    int t=0;
    for(int l=1,r=deep[0];l<r;) {
        if(deep[l]+deep[r]<=k) 
            t+=r-l,l++;
        else r--;
    }
    return t;
}
void work(int x) {
    ans+=calc(x,0);
    vis[x]=1;
    for(int i=head[x];i;i=edge[i].next) {
        int v=edge[i].v;
        if(vis[v])continue;
        ans-=calc(v,edge[i].w);
        sum=son[v];
        root=0;
        getroot(v,root);
        work(root);
    }
}
int main() {
    while(11101001) {
        ans=0;num=0;root=0;
        n=read(),k=read();
        if(!n&&!k)break;
        std::memset(vis,0,sizeof vis);
        std::memset(head,0,sizeof head);
        for(int a,b,c,i=1;i<n;++i) {
            a=read(),b=read(),c=read();
            add_edge(a,b,c);
            add_edge(b,a,c);
        }
        sum=n;f[0]=INF;
        getroot(1,0);
        work(root);
        printf("%d
",ans);
    }
    return 0;
}