BZOJ 1468: Tree

点分治秒了这题

#include<cstdio>#include<algorithm>#define RG register#define N 40001using namespace std;struct edge{int y,v,next;}e[N*2];int head[N];inline int read(){int a=0,f=1;static char c=getchar();while(c<'0'||c>'9'){if(c=='-')f=-1;c=getchar();}while(c>='0'&&c<='9'){a=a*10+c-'0';c=getchar();}return a*f;}int n,k,ans,root,cnt,sum;int mark[N],son[N],mx[N],deep[N],d[N];inline void add(int x,int y,int v){e[++cnt].next=head[x];head[x]=cnt;e[cnt].y=y,e[cnt].v=v;}void getr(int x,int fa){son[x]=1;mx[x]=0;for(RG int i=head[x];i;i=e[i].next){if(mark[e[i].y]||e[i].y==fa) continue;getr(e[i].y,x);son[x]+=son[e[i].y];mx[x]=max(mx[x],son[e[i].y]);}mx[x]=max(mx[x],sum-son[x]);if(mx[x]<mx[root]) root=x;}void getd(int x,int fa){deep[++deep[0]]=d[x];for(RG int i=head[x];i;i=e[i].next){if(mark[e[i].y]||e[i].y==fa) continue;d[e[i].y]=d[x]+e[i].v;getd(e[i].y,x);}}int cal(int x,int v){d[x]=v;deep[0]=0;getd(x,0);sort(deep+1,deep+deep[0]+1);int l=1,r=deep[0],sum=0;while(l<r){if(deep[l]+deep[r]<=k) {sum+=r-l;++l;}else --r;}return sum;}void solve(int x){ans+=cal(x,0);mark[x]=1;for(RG int i=head[x];i;i=e[i].next){if(mark[e[i].y]) continue;ans-=cal(e[i].y,e[i].v);sum=son[e[i].y];root=0;getr(e[i].y,0);solve(root);}}int main(){n=read();RG int i,x,y,z;ans=root=cnt=0;for(i=1;i<n;++i,x=read(),y=read(),z=read(),add(x,y,z),add(y,x,z));k=read();mx[0]=2e8;sum=n;getr(1,0);solve(root);printf("%d\n",ans);return 0;}