[poj1741]Tree-蒲公英云

[poj1741]Tree

点分治模板题，可以将同一棵树的链分为两种：1.通过重心；2.在子树内部。第2种可以搜下去，第1种的答案即$\sum_{i,j}[di+dj<=m]-\sum\limits_{i,j在同一个子树内}[di+dj<=m]$。

1 #include<cstdio>
 2 #include<cstring>
 3 #include<algorithm>
 4 using namespace std;
 5 #define N 10005
 6 struct ji{
 7     int nex,to,len;
 8 }edge[N<<1];
 9 int E,r,n,m,x,y,z,ans,a[N],head[N],vis[N],sz[N];
10 void add(int x,int y,int z){
11     edge[E].nex=head[x];
12     edge[E].to=y;
13     edge[E].len=z;
14     head[x]=E++;
15 }
16 void tot(int k,int fa,int sh){
17     a[++a[0]]=sh;
18     for(int i=head[k];i!=-1;i=edge[i].nex)
19         if ((!vis[edge[i].to])&&(edge[i].to!=fa))
20             tot(edge[i].to,k,sh+edge[i].len);
21 }
22 int calc(int k,int p){
23     int ans=a[0]=0;
24     tot(k,0,p);
25     sort(a+1,a+a[0]+1);
26     for(int i=1,j=a[0];i<j;i++){
27         while ((i<j)&&(a[i]+a[j]>m))j--;
28         ans+=j-i;
29     }
30     return ans;
31 }
32 void get(int k,int fa){
33     int ma=0;
34     sz[k]=1;
35     for(int i=head[k];i!=-1;i=edge[i].nex)
36         if ((!vis[edge[i].to])&&(edge[i].to!=fa)){
37             get(edge[i].to,k);
38             sz[k]+=sz[edge[i].to];
39             ma=max(ma,sz[edge[i].to]);
40         }
41     ma=max(ma,sz[0]-sz[k]);
42     if (ma<=sz[0]/2)r=k;
43 }
44 void dfs(int k){
45     ans+=calc(k,0);
46     vis[k]=1;
47     get(k,0);
48     for(int i=head[k];i!=-1;i=edge[i].nex)
49         if (!vis[edge[i].to]){
50             ans-=calc(edge[i].to,edge[i].len);
51             sz[0]=sz[edge[i].to];
52             get(edge[i].to,0);
53             dfs(r);
54         }
55 }
56 int main(){
57     while (scanf("%d%d",&n,&m)!=EOF){
58         if ((!n)&&(!m))return 0;
59         E=ans=0;
60         memset(head,-1,sizeof(head));
61         memset(vis,0,sizeof(vis));
62         for(int i=1;i<n;i++){
63             scanf("%d%d%d",&x,&y,&z);
64             add(x,y,z);
65             add(y,x,z);
66         }
67         sz[0]=n;
68         get(1,0);
69         dfs(r);
70         printf("%d\n",ans);
71     }
72 }