BZOJ 1468 Tree 【模板】树上点分治

 1 #include<cstdio>
 2 #include<algorithm>
 3 #define N 50010
 4 #define M 500010
 5 #define rg register
 6 #define LL long long
 7 using namespace std;
 8 int n,m,cnt,root,tot,totnow,totbf;
 9 int last[N],size[N],mxsize[N],dep[N],tmp[N],mg[N];
10 LL ans=0;
11 bool v[N];
12 struct edge{
13     int to,pre,dis;
14 }e[M];
15 inline int read(){
16     int k=0,f=1; char c=getchar();
17     while(c<‘0‘||c>‘9‘)c==‘-‘&&(f=-1),c=getchar();
18     while(‘0‘<=c&&c<=‘9‘)k=k*10+c-‘0‘,c=getchar();
19     return k*f;
20 }
21 void getroot(int x,int fa){
22     size[x]=1; mxsize[x]=0;
23     for(rg int i=last[x],to;i;i=e[i].pre)if(!v[to=e[i].to]&&to!=fa){
24         getroot(to,x); size[x]+=size[to];
25         mxsize[x]=max(mxsize[x],size[to]);
26     }
27     mxsize[x]=max(mxsize[x],cnt-mxsize[x]);
28     if(!root||mxsize[x]<mxsize[root]) root=x;
29 }
30 void getdep(int x,int fa,int d){
31     dep[++totnow]=d;
32     for(rg int i=last[x],to;i;i=e[i].pre)
33         if(!v[to=e[i].to]&&to!=fa) getdep(to,x,d+e[i].dis);
34 }
35 void dfs(int x){
36     v[x]=1; totbf=1; tmp[1]=0;
37     for(rg int i=last[x],to;i;i=e[i].pre)if(!v[to=e[i].to]){
38         totnow=0; getdep(to,x,e[i].dis); sort(dep+1,dep+1+totnow);
39         int p=1;
40         for(rg int j=totbf;j;j--){
41             while(p<=totnow&&tmp[j]+dep[p]<=m) p++;
42             ans+=p-1;
43         }
44         p=1; int p2=1,totmg=0;
45         while(p<=totnow&&p2<=totbf) 
46             mg[++totmg]=dep[p]<tmp[p2]?dep[p++]:tmp[p2++];
47         while(p<=totnow) mg[++totmg]=dep[p++];
48         while(p2<=totbf) mg[++totmg]=tmp[p2++];
49         for(rg int j=1;j<=totmg;j++) tmp[j]=mg[j];
50         totbf=totmg;
51     }
52     for(rg int i=last[x],to;i;i=e[i].pre)if(!v[to=e[i].to]){
53         root=0; cnt=size[to];
54         getroot(to,x); dfs(root);
55     }
56 }
57 int main(){
58     n=read();
59     for(rg int i=1;i<n;i++){
60         int u=read(),v=read(),w=read();
61         e[++tot]=(edge){v,last[u],w}; last[u]=tot;
62         e[++tot]=(edge){u,last[v],w}; last[v]=tot;
63     }
64     m=read();
65     cnt=n; getroot(1,0); dfs(root);
66     printf("%lld\n",ans);
67 }