POJ1655 【树的分治】

楼教主的题。

//每次选择树的重心，删去重心形成多棵子树（子树一定小于上一层的一半），最多log(n)层 #include<cstdio> #include<cstring> #include<algorithm> using namespace std;  const int N=22222;  int ev[N],ew[N],nxt[N],head[N],e; int vis[N],dp[N],arr[N]; int n,k,core,sz,pmn,ed;  void init() {     memset(head,-1,sizeof(head));     memset(vis,0,sizeof(vis));     e=ed=0; } void add(int u,int v,int w) {     ev[e]=v,ew[e]=w,nxt[e]=head[u];head[u]=e++; }  void calsz(int u,int p)//辅助求重心 {     for(int i=head[u];~i;i=nxt[i]) if(ev[i]!=p&&!vis[ev[i]]) calsz(ev[i],u),sz++; } void calcore(int u,int p)//求重心 {     dp[u]=1;int mx=0,v;     for(int i=head[u];~i;i=nxt[i]) if(ev[i]!=p&&!vis[v=ev[i]])     {         calcore(v,u);         dp[u]+=dp[v];         mx=max(mx,dp[v]);     }     mx=max(mx,sz-dp[u]);     if(mx<pmn) pmn=mx,core=u; } int cntnum(int *arr,int len)//arr[0]-arr[len-1]满足条件的点对数 {     int ans=0;     sort(arr,arr+len);     for(int i=0,j=len-1;i<j;i++)     {         while(i<j&&arr[i]+arr[j]>k) j--;         ans+=j-i;     }     return ans; } void make(int u,int p,int len)//将子树所有结点到当前根的距离放入arr数组中 {     arr[ed++]=len;     for(int i=head[u];~i;i=nxt[i]) if(ev[i]!=p&&!vis[ev[i]]) make(ev[i],u,len+ew[i]); } int dfs(int u) {     pmn=1e8; sz=1;     calsz(u,u);     calcore(u,u);     ed=vis[u=core]=1;     int ans=0;     for(int i=head[u];~i;i=nxt[i])     {         int st=ed;         if(!vis[ev[i]]) make(ev[i],u,ew[i]);         ans-=cntnum(arr+st,ed-st);//先将下一层子树内部满足条件的点对减掉     }     ans+=cntnum(arr,ed);     for(int i=head[u];~i;i=nxt[i]) if(!vis[ev[i]]) ans+=dfs(ev[i]);//统计当前子树点对     return ans; }  int main() {     while(scanf("%d%d",&n,&k),n||k) {         init();         for(int i=1; i<n; i++) {             int u,v,w;             scanf("%d%d%d",&u,&v,&w);             add(u,v,w),add(v,u,w);         }         printf("%d\n",dfs(1));     }     return 0; }