poj 1741（树的点分治）

  1 #include<iostream>
  2 #include<cstdio>
  3 #include<cmath>
  4 #include<cstring>
  5 #include<algorithm>
  6 #include<queue>
  7 #include<map>
  8 #include<set>
  9 #include<vector>
 10 #include<cstdlib>
 11 #include<stack>
 12 #include<string>
 13 #define eps 0.000000001
 14 typedef long long ll;
 15 typedef unsigned long long LL;
 16 using namespace std;
 17 const int INF=0x3f3f3f3f;
 18 const int N=100000+10;
 19 int son[N];
 20 int ans;
 21 int num;
 22 int p;
 23 int head[N];
 24 int vis[N];
 25 int minn;
 26 int root;
 27 int k;
 28 int dis[N];
 29 void init(){
 30     num=0;
 31     memset(vis,0,sizeof(vis));
 32     memset(head,-1,sizeof(head));
 33 }
 34 struct node{
 35     int to,next,w;
 36 }edge[N];
 37 void add(int u,int v,int w){
 38     edge[num].to=v;
 39     edge[num].w=w;
 40     edge[num].next=head[u];
 41     head[u]=num++;
 42 }
 43 void getroot(int x,int y,int n){
 44     son[x]=1;
 45     int maxx=0;
 46     for(int i=head[x];i!=-1;i=edge[i].next){
 47         int v=edge[i].to;
 48         if(v==y||vis[v])continue;
 49         getroot(v,x,n);
 50         son[x]=son[x]+son[v];
 51         maxx=max(maxx,son[v]);
 52     }
 53     maxx=max(maxx,n-son[x]);
 54     //cout<<maxx<<" "<<x<<endl;
 55     if(minn>maxx){
 56         minn=maxx;
 57         root=x;
 58     }
 59 }
 60 void getdeep(int x,int v0,int y){
 61     dis[p++]=y;
 62     for(int i=head[x];i!=-1;i=edge[i].next){
 63         int v=edge[i].to;
 64         if(v==v0||vis[v])continue;
 65         getdeep(v,x,y+edge[i].w);
 66     }
 67 }
 68 int cal(int x,int y){
 69     int res=0;
 70     p=0;
 71     getdeep(x,0,y);
 72     sort(dis,dis+p);
 73     int l=0;
 74     int r=p-1;
 75     while(l<r){
 76         if(dis[l]+dis[r]<=k){
 77             res=res+(r-l);
 78             l++;
 79         }
 80         else{
 81             r--;
 82         }
 83     }
 84     return res;
 85 }
 86 void work(int x){
 87     ans=ans+cal(x,0);
 88     //cout<<cal(x,0)<<" ";
 89    // cout<<endl;
 90     vis[x]=1;
 91     for(int i=head[x];i!=-1;i=edge[i].next){
 92         int v=edge[i].to;
 93         if(vis[v]==0){
 94             ans=ans-cal(v,edge[i].w);
 95             minn=INF;
 96             getroot(v,0,son[v]);
 97             work(root);
 98         }
 99     }
100 }
101 int main(){
102     int n;
103     while(scanf("%d%d",&n,&k)!=EOF){
104         if(n==0&&k==0)break;
105         init();
106         for(int i=0;i<n-1;i++){
107             int u,v,w;
108             scanf("%d%d%d",&u,&v,&w);
109             add(u,v,w);
110             add(v,u,w);
111         }
112         ans=0;
113         minn=INF;
114         getroot(1,-1,n);
115         //for(int i=1;i<=n;i++)cout<<son[i]<<" ";
116        // cout<<endl;
117        //cout<<root<<" "<<minn<<endl;
118         work(root);
119         cout<<ans<<endl;
120     }
121     return 0;
122 }