POJ 1741/1987 树的点分治

树的点分治，主要思想是每次找子树的重心，计算经过根节点的情况数，再减去点对属于同一子树的情况。

 #include <iostream>
 #include <vector>
 #include <algorithm>
 #include <string>
 #include <string.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <queue>
 #include <stack>
 #include <map>
 #include <set>
 #include <cmath>
 #include <ctime>
 #include <cassert>
 #include <sstream>
 using namespace std;

 const int N=;
 struct Edge {
     int to,next;
     int w;
     Edge(){}
     Edge (int _t,int _n,int _w=) {
         to=_t;
         next=_n;
         w=_w;
     }
 }edge[N<<];
 int idx,head[N];
 void addEdge (int u,int v,int w) {
     ++idx;
     edge[idx]=Edge(v,head[u],w);
     head[u]=idx;
 }
 bool vis[N];

 int K;// 输入中的k

 int dis[N],disCnt;
 void getDis(int u,int f,int d) {
     dis[disCnt++]=d;
     for (int k=head[u];k;k=edge[k].next) {
         int v=edge[k].to;
         if (vis[v]||v==f) continue;
         getDis(v,u,d+edge[k].w);
     }
 }

 int bal,cmp;
 int getBal(int u,int f) {
     int son=;
     int dp=;
     for (int k=head[u];k;k=edge[k].next) {
         int v=edge[k].to;
         if (vis[v]||v==f) continue;
         int subSon=getBal(v,u);
         son+=subSon;
         dp=max(dp,subSon);
     }
     dp=max(dp,disCnt-son);
     if (dp<cmp){
         cmp=dp;
         bal=u;
     }
     return son;
 }

 int calc(int u,int initDis) {
     disCnt=;
     getDis(u,-,initDis);
     sort(dis,dis+disCnt);
     int ret=;
     int l=,r=disCnt-;
     while (l<r) {
         if (dis[l]+dis[r]>K) r--;
         else ret+=r-l++;
     }
     return ret;
 }

 int ans=;
 void solve(int u) {
     cmp=~0U>>;
     getBal(u,-);
     ans+=calc(bal,);
     vis[bal]=true;
     for (int k=head[bal];k;k=edge[k].next) {
         int v=edge[k].to;
         if (vis[v]) continue;
         ans-=calc(v,edge[k].w);// 减去子树内的重复计算的情况
         solve(v);
     }
 }

 void init(int n) {
     idx=;memset(head,,sizeof head);
     memset(vis,false,sizeof vis);
     disCnt=n;// 由于在计算重心时，可以直接用之前算dis时的disCnt，所以这里初始化为n
     ans=;
 }
 int main () {
     int n;
     while (scanf("%d %d",&n,&K)!=EOF) {
         if (n==&&K==) break;
         init(n);
         for (int i=;i<n;i++) {
             int u,v,w;
             scanf("%d %d %d",&u,&v,&w);
             addEdge(u,v,w);
             addEdge(v,u,w);
         }
         solve();
         printf("%d\n",ans);
     }
     return ;
 }