【Luogu 4178】Tree

 1 #include<iostream>
 2 #include<cstdio>
 3 using namespace std;
 4 const int M=100009;
 5 int n,k,num=0,minn,rt,size,ans=0,tot,top;
 6 int head[M],siz[M],p[M],g[M],b[M];
 7 struct P{int to,ne,w;}e[M<<1];
 8 bool vis[M];
 9 int read(){
10     int rex=0,f=1;char ch=getchar();
11     while(ch<‘0‘||ch>‘9‘){if(ch==‘-‘)f=-1;ch=getchar();}
12     while(ch>=‘0‘&&ch<=‘9‘){rex=rex*10+ch-‘0‘;ch=getchar();}
13     return rex*f;
14 }
15 void getrt(int u,int fa){
16     siz[u]=1;int ma=0;
17     for(int i=head[u];i;i=e[i].ne){
18         int v=e[i].to;if(v==fa||vis[v])continue;
19         getrt(v,u);siz[u]+=siz[v];ma=max(ma,siz[v]);
20     }
21     ma=max(ma,size-siz[u]);
22     if(minn>ma){minn=ma,rt=u;}
23 }
24 void dfs(int u,int fa,int w){
25     b[++tot]=w;
26     for(int i=head[u];i;i=e[i].ne){
27         int v=e[i].to;if(v==fa||vis[v])continue;
28         dfs(v,u,w+e[i].w);
29     }
30 }
31 void add(int x,int y){
32     for(;x<=k;x+=x&-x)g[x]+=y;
33 }
34 int sum(int x,int rex=0){
35     for(;x;x-=x&-x)rex+=g[x];
36     return rex;
37 }
38 void get(int v,int u,int w){
39     tot=0;
40     dfs(v,u,w);
41     for(int i=1;i<=tot;++i){
42         if(b[i]<=k)ans++;
43         if(b[i]>=k)continue;
44         p[++top]=b[i];
45         ans+=sum(k-b[i]);
46     }
47     for(int i=1;i<=tot;++i){
48         if(b[i]>=k)continue;
49         add(b[i],1);
50     }
51 }
52 void work(int u){
53     vis[u]=1;top=0;
54     for(int i=head[u];i;i=e[i].ne){
55         int v=e[i].to;
56         if(!vis[v])get(v,u,e[i].w);
57     }
58     for(int i=1;i<=top;++i)add(p[i],-1);
59     for(int i=head[u];i;i=e[i].ne){
60         int v=e[i].to;if(vis[v])continue;
61         minn=1e9;size=siz[u];
62         getrt(v,0);
63         work(rt);
64     }
65 }
66 int main(){
67     n=read();
68     for(int i=1,u,v,w;i<n;++i){
69         u=read(),v=read(),w=read();
70         e[++num]=(P){v,head[u],w};head[u]=num;
71         e[++num]=(P){u,head[v],w};head[v]=num;
72     }
73     k=read();
74     minn=1e9;size=n;
75     getrt(1,0);
76     work(rt);
77     printf("%d",ans);
78     return 0;
79 }