POJ 1741/1987 树的点分治

  1 #include <iostream>
  2 #include <vector>
  3 #include <algorithm>
  4 #include <string>
  5 #include <string.h>
  6 #include <stdio.h>
  7 #include <stdlib.h>
  8 #include <queue>
  9 #include <stack>
 10 #include <map>
 11 #include <set>
 12 #include <cmath>
 13 #include <ctime>
 14 #include <cassert>
 15 #include <sstream>
 16 using namespace std;
 17 
 18 const int N=10010;
 19 struct Edge {
 20     int to,next;
 21     int w;
 22     Edge(){}
 23     Edge (int _t,int _n,int _w=0) {
 24         to=_t;
 25         next=_n;
 26         w=_w;
 27     }
 28 }edge[N<<1];
 29 int idx,head[N];
 30 void addEdge (int u,int v,int w) {
 31     ++idx;
 32     edge[idx]=Edge(v,head[u],w);
 33     head[u]=idx;
 34 }
 35 bool vis[N];
 36 
 37 int K;// 输入中的k
 38 
 39 int dis[N],disCnt;
 40 void getDis(int u,int f,int d) {
 41     dis[disCnt++]=d;
 42     for (int k=head[u];k;k=edge[k].next) {
 43         int v=edge[k].to;
 44         if (vis[v]||v==f) continue;
 45         getDis(v,u,d+edge[k].w);
 46     }
 47 }
 48 
 49 int bal,cmp;
 50 int getBal(int u,int f) {
 51     int son=1;
 52     int dp=0;
 53     for (int k=head[u];k;k=edge[k].next) {
 54         int v=edge[k].to;
 55         if (vis[v]||v==f) continue;
 56         int subSon=getBal(v,u);
 57         son+=subSon;
 58         dp=max(dp,subSon);
 59     }
 60     dp=max(dp,disCnt-son);
 61     if (dp<cmp){
 62         cmp=dp;
 63         bal=u;
 64     }
 65     return son;
 66 }
 67 
 68 int calc(int u,int initDis) {
 69     disCnt=0;
 70     getDis(u,-1,initDis);
 71     sort(dis,dis+disCnt);
 72     int ret=0;
 73     int l=0,r=disCnt-1;
 74     while (l<r) {
 75         if (dis[l]+dis[r]>K) r--;
 76         else ret+=r-l++;
 77     }
 78     return ret;
 79 }
 80 
 81 int ans=0;
 82 void solve(int u) {
 83     cmp=~0U>>1;
 84     getBal(u,-1);
 85     ans+=calc(bal,0);
 86     vis[bal]=true;
 87     for (int k=head[bal];k;k=edge[k].next) {
 88         int v=edge[k].to;
 89         if (vis[v]) continue;
 90         ans-=calc(v,edge[k].w);// 减去子树内的重复计算的情况
 91         solve(v);
 92     }
 93 }
 94 
 95 void init(int n) {
 96     idx=1;memset(head,0,sizeof head);
 97     memset(vis,false,sizeof vis);
 98     disCnt=n;// 由于在计算重心时，可以直接用之前算dis时的disCnt，所以这里初始化为n
 99     ans=0;
100 }
101 int main () {
102     int n;
103     while (scanf("%d %d",&n,&K)!=EOF) {
104         if (n==0&&K==0) break;
105         init(n);
106         for (int i=1;i<n;i++) {
107             int u,v,w;
108             scanf("%d %d %d",&u,&v,&w);
109             addEdge(u,v,w);
110             addEdge(v,u,w);
111         }
112         solve(1);
113         printf("%d\n",ans);
114     }
115     return 0;
116 }