[BZOJ1468]Tree 点分治

#include<iostream>
#include<cstdio>
#include<algorithm>
using namespace std;
const int maxn=40010;
int n,k,top,ans;
bool vis[maxn];
int st[maxn],sz[maxn];
struct edge
{
int t,w;
edge *next;
}*con[maxn];
void ins(int x,int y,int w)
{
edge *p=new edge;
p->t=y;
p->w=w;
p->next=con[x];
con[x]=p;
}
void getroot(int v,int fa,int size,int &root)
{
bool flag=1;
sz[v]=1;
for(edge *p=con[v];p;p=p->next)
if(p->t!=fa&&!vis[p->t])
{
getroot(p->t,v,size,root);
sz[v]+=sz[p->t];
if(sz[p->t]*2>size) flag=0;
}
if(size-sz[v]*2>0) flag=0;
if(flag) root=v;
}
void dfs(int v,int fa,int d)
{
st[++top]=d;
for(edge *p=con[v];p;p=p->next)
if(!vis[p->t]&&p->t!=fa) dfs(p->t,v,d+p->w);
}
void cal(int v,int f,int d)
{
top=0;
dfs(v,-1,d);
sort(st+1,st+top+1);
for(int hd=1,tl=top;hd<tl;hd++)
{
while(hd<tl&&st[hd]+st[tl]>k) tl--;
ans+=(tl-hd)*f;
}
}
void solve(int v,int size)
{
int rt;
getroot(v,-1,size,rt);
vis[rt]=1;
cal(rt,1,0);
for(edge *p=con[rt];p;p=p->next)if(!vis[p->t]) cal(p->t,-1,p->w);
for(edge *p=con[rt];p;p=p->next)if(!vis[p->t]) {solve(p->t,(sz[p->t]>sz[rt]?size-sz[rt]:sz[p->t]));}
}
int main()
{
scanf("%d",&n);
for(int i=1;i<n;i++)
{
int x,y,w;
scanf("%d%d%d",&x,&y,&w);
ins(x,y,w);
ins(y,x,w);
}
scanf("%d",&k);
ans=0;
solve(1,n);
printf("%d",ans);
return 0;
}