poj 1741 Tree 树上的分治+求树的重心

//poj 1741
//sep9
#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;
const int MAXN=10024;
int ans,n,K,e;
int subtree_size[MAXN],head[MAXN];
bool centroid[MAXN];
struct EDGE
{
int v,w,nxt;
}edge[2*MAXN];

int compute_subtree_size(int u,int p)
{
int c=1;
for(int i=head[u];i!=-1;i=edge[i].nxt){
int v=edge[i].v,w=edge[i].w;
if(v==p||centroid[v])
continue;
c+=compute_subtree_size(v,u);
}
return subtree_size[u]=c;
}

pair<int,int> search_centroid(int u,int p,int t)
{
pair<int,int> res=make_pair(INT_MAX,-1);
int s=1,m=0;
for(int i=head[u];i!=-1;i=edge[i].nxt){
int v=edge[i].v,w=edge[i].w;
if(v==p||centroid[v]) continue;

res=min(res,search_centroid(v,u,t));

m=max(m,subtree_size[v]);
s+=subtree_size[v];
}
m=max(m,t-s);
res=min(res,make_pair(m,u));
return res;
}

void enumerate_paths(int u,int p,int len,vector<int> &ds)
{
ds.push_back(len);
for(int i=head[u];i!=-1;i=edge[i].nxt){
int v=edge[i].v,w=edge[i].w;
if(v==p||centroid[v]) continue;
enumerate_paths(v,u,w+len,ds);
}
}

int count_pairs(vector<int> &ds)
{
int res=0;
sort(ds.begin(),ds.end());
int j=ds.size()-1;
for(int i=0;i<ds.size();++i){
while(j>=0&&ds[i]+ds[j]>K) --j;
if(j>=0) res+=j+1-(j>=i?1:0);
}
return res/2;
}

void solve_subproblem(int u)
{
compute_subtree_size(u,-1);
int s=search_centroid(u,-1,subtree_size[u]).second;
centroid[s]=true;
for(int i=head[s];i!=-1;i=edge[i].nxt){
if(centroid[edge[i].v]) continue;
solve_subproblem(edge[i].v);
}
vector<int> ds;
ds.push_back(0);
for(int i=head[s];i!=-1;i=edge[i].nxt){
if(centroid[edge[i].v]) continue;
vector<int> tds;
enumerate_paths(edge[i].v,s,edge[i].w,tds);

ans-=count_pairs(tds);
ds.insert(ds.end(),tds.begin(),tds.end());
}
ans+=count_pairs(ds);
centroid[s]=false;//important
}

void solve()
{
ans=0;
memset(centroid,false,sizeof(centroid));
solve_subproblem(1);
printf("%d\n",ans);
}

int main()
{
while(scanf("%d%d",&n,&K)==2&&!(n==0&&K==0)){
e=0;
memset(head,-1,sizeof(head));
for(int i=1;i<n;++i){
int u,v,w;
scanf("%d%d%d",&u,&v,&w);
edge[e].v=v,edge[e].w=w,edge[e].nxt=head[u],head[u]=e++;
edge[e].v=u,edge[e].w=w,edge[e].nxt=head[v],head[v]=e++;
}
solve();
}
return 0;
}