【模板】点分治2

  1 // Cease to struggle and you cease to live
  2 // p4178
  3 #include <iostream>
  4 #include <cmath>
  5 #include <cstdio>
  6 #include <cstring>
  7 #include <algorithm>
  8 #include <queue>
  9 #include <vector>
 10 #include <set>
 11 #include <map>
 12 #include <stack>
 13 using namespace std;
 14 typedef long long ll;
 15 const int N=4e4+8;
 16 const int maxn=20000+7;
 17 int n,k,cnt;
 18 ll ans;
 19 ll tr[maxn<<4];
 20 int sav[N],era[N];
 21 int h[N],siz[N],dis[N],all,root,mx;
 22 bool vis[N];
 23 struct edge{
 24     int to,nex,w;
 25 }e[N<<1];
 26 void add(int u,int v,int w){
 27     e[++cnt]=(edge){v,h[u],w};
 28     h[u]=cnt;
 29 }
 30 void addval(int x,ll v){
 31     for(;x<=k;x+=x&(-x)) tr[x]+=v;
 32 }
 33 ll sum(int x){
 34     ll res=0;
 35     for(;x>0;x-=x&(-x)) res+=tr[x];
 36     return res;
 37 }
 38 void getrt(int u,int fa){
 39     siz[u]=1;
 40     int num=0;
 41     for(int i=h[u];i;i=e[i].nex){
 42         int v=e[i].to;
 43         if(v==fa || vis[v]) continue;
 44         getrt(v,u);
 45         siz[u]+=siz[v];
 46         num=max(num,siz[v]);
 47     }
 48     num=max(num,all-siz[u]);
 49     if(num<mx){
 50         mx=num;
 51         root=u;
 52     }
 53 }
 54 void getdis(int u,int fa){
 55     if(dis[u]>k) return;
 56     sav[++sav[0]]=era[++era[0]]=dis[u];
 57     for(int i=h[u];i;i=e[i].nex){
 58         int v=e[i].to;
 59         if(vis[v] || v==fa) continue;
 60         dis[v]=dis[u]+e[i].w;
 61         getdis(v,u);
 62     }
 63 }
 64 void cal(int u){
 65     era[0]=0;
 66     for(int i=h[u];i;i=e[i].nex){
 67         int v=e[i].to;
 68         if(vis[v]) continue;
 69         sav[0]=0;
 70         dis[v]=e[i].w;
 71         getdis(v,u);
 72         ans+=sav[0];
 73         for(int j=1;j<=sav[0];++j){
 74             ans+=sum(k-sav[j]);
 75         }
 76         for(int j=1;j<=sav[0];++j){
 77             addval(sav[j],1);
 78         }
 79     }
 80     for(int i=1;i<=era[0];++i){
 81         if(era[i]>k) continue;
 82         addval(era[i],-1);
 83     }
 84 }
 85 void dfs(int u){
 86     cal(u);
 87     vis[u]=1;
 88     for(int i=h[u];i;i=e[i].nex){
 89         int v=e[i].to;
 90         if(vis[v]) continue;
 91         all=siz[v];
 92         mx=0x3f3f3f3f;
 93         getrt(v,0);
 94         dfs(root);
 95     }
 96 }
 97 int main() {
 98     scanf("%d",&n);
 99     ans=cnt=0;
100     for(int i=1;i<n;++i){
101         int a,b,c;
102         scanf("%d%d%d",&a,&b,&c);
103         add(a,b,c);add(b,a,c);
104     }
105     scanf("%d",&k);
106     all=n;
107     mx=0x3f3f3f3f;
108     getrt(1,0);
109     dfs(root);
110     printf("%lld\n",ans);
111     return 0;
112 }

View Code