树的点分治——模板

Posted 2020-10-08 Troy Ricardo

  1 #define Troy 9/28/2017
  2 
  3 #define inf 0x7fffffff
  4 
  5 #include "cstdio"
  6 #include "cstring"
  7 #include "algorithm"
  8 using std::sort;
  9 
 10 typedef long long ll;
 11 ll ans;
 12 
 13 inline int max(int a,int b){return a>b?a:b;}
 14 inline int min(int a,int b){return a<b?a:b;}
 15 inline int read(){
 16     int s=0,k=1;char ch=getchar();
 17     while(ch<‘0‘|ch>‘9‘)    ch==‘-‘?k=-1:0,ch=getchar();
 18     while(ch>47&ch<=‘9‘)    s=s*10+(ch^48),ch=getchar();
 19     return s*k;
 20 }
 21 const int N=10050;
 22 
 23 struct edges {
 24     int v,w;edges *last;
 25 }edge[N<<1],*head[N];int cnt;
 26 inline void push(int u,int v,int w){
 27     edge[++cnt]=(edges){v,w,head[u]};head[u]=edge+cnt;
 28     edge[++cnt]=(edges){u,w,head[v]};head[v]=edge+cnt;
 29 }
 30 bool del[N];int n,K;
 31 int top;bool vis[N];
 32 
 33 inline void clear(){
 34     cnt=0;ans=0;
 35     memset(vis,0,sizeof(vis));
 36     memset(head,0,sizeof(head));
 37 }
 38 
 39 int size[N],heavy[N],root,num,dis[N];
 40 
 41 inline void dfs_size(int x,int fa){
 42     size[x]=1;
 43     heavy[x]=0;
 44     for(edges *i=head[x];i;i=i->last)   if(vis[i->v]||i->v==fa)  continue;
 45     else{
 46         dfs_size(i->v,x);
 47         heavy[x]=max(size[i->v],heavy[x]);
 48         size[x]+=size[i->v];
 49     }
 50 }
 51 
 52 inline void find_root(int tot,int x,int fa){
 53     heavy[x]=max(tot-size[x],heavy[x]);
 54     if(heavy[x]<top)    top=heavy[x],root=x;
 55     for(edges *i=head[x];i;i=i->last){
 56         if(i->v==fa|vis[i->v])  continue;
 57         find_root(tot,i->v,x);
 58     }
 59 }
 60 
 61 inline void dfs_dis(int x,int d,int fa){
 62     dis[num++]=d;
 63     for(edges *i=head[x];i;i=i->last){
 64         if(i->v==fa|vis[i->v])    continue;
 65         dfs_dis(i->v,d+i->w,x);
 66     }
 67 }
 68 
 69 inline ll calc(int x,int d){
 70     ll ret=0;
 71     num=0;
 72     dfs_dis(x,d,x);
 73     sort(dis,dis+num);
 74     int l=0,r=num-1;
 75     while(l<r){
 76         while(l<r&&dis[l]+dis[r]>K) r--;
 77         ret+=r-l;
 78         l++;
 79     }return ret;
 80 }
 81 
 82 inline void solve(int u){
 83     top=inf;
 84     dfs_size(u,u);
 85     find_root(size[u],u,u);
 86     ans+=calc(root,0);
 87     vis[root]=true;
 88     for(edges *i=head[root];i;i=i->last)
 89         if(!vis[i->v]){
 90 
 91             ans-=calc(i->v,i->w);
 92             solve(i->v);
 93         }
 94 }
 95 
 96 int main(){
 97 #ifdef Troy
 98     while(scanf("%d%d",&n,&K)==2&&n&&K){
 99         clear();
100         for(int i=1,u,v,w;i<n;i++){
101             u=read(),v=read(),w=read();
102             push(u,v,w);
103         }
104         solve(1);
105         printf("%lld\n",ans);
106     }
107 #endif
108 }
109 /*
110 5 4
111 1 2 3
112 1 3 1
113 1 4 2
114 3 5 1
115 0 0
116 */