poj 1741 Tree 树的分治

#include<iostream>
#include<cmath>
#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
#define NMAX 10010
struct node
{
int num,weight;
node *next;
};
node *graph[NMAX];
node memo[2*NMAX];
bool use[NMAX];
int size[NMAX],d[NMAX];
int f[NMAX];
int n,m,top,tot,root;
int sumsize,minsize;
void add(int x,int y,int z)
{
node *p=&memo[top++];
p->num=y; p->weight=z; p->next=graph[x]; graph[x]=p;
p=&memo[top++];
p->num=x; p->weight=z; p->next=graph[y]; graph[y]=p;
}
void get_root(int i,int fa)
{
int big=-1;
size[i]=1;
for(node *p=graph[i];p;p=p->next)
if(!use[p->num]&&p->num!=fa)
{
get_root(p->num,i);
size[i]+=size[p->num];
if(size[p->num]>big) big=size[p->num];
}
if(sumsize-size[i]>big) big=sumsize-size[i];
if(big<minsize) minsize=big,root=i;
}
void get_dis(int i,int dis,int fa)
{
d[++tot]=dis;
for(node *p=graph[i];p;p=p->next)
if(!use[p->num]&&p->num!=fa)
get_dis(p->num,dis+p->weight,i);
}

int count(int x,int dis)
{
int s=0;
tot=0;
get_dis(x,dis,-1);
sort(d+1,d+tot+1);
for(int i=1,j=tot;i<=j;i++)
{
while(d[i]+d[j]>m&&i<=j) j--;
if(i<j) s+=j-i;
}
return s;
}

void dfs(int i)
{
sumsize=minsize=size[i];
get_root(i,-1);
i=root;
f[i]=count(i,0);
//cout<<i<<" w "<<f[i]<<endl;
use[i]=1;
for(node *p=graph[i];p;p=p->next)
if(!use[p->num])
{
f[i]-=count(p->num,p->weight);
dfs(p->num);
}
}
int main()
{
int i;
int x,y,z;
while(scanf("%d%d",&n,&m)!=EOF&&n)
{
int ans=0;
top=0;
memset(graph,0,sizeof(graph));
memset(use,0,sizeof(use));
for(i=1;i<n;i++)
{
scanf("%d%d%d",&x,&y,&z);
add(x,y,z);
}
size[1]=n;
dfs(1);
for(i=1;i<=n;i++) ans+=f[i];
printf("%d\n",ans);
}
return 0;
}