Codeforces 161D Distance in Tree（树的点分治）

 1 #include<cstdio>
 2 #include<cstring>
 3 #include<algorithm>
 4 using namespace std;
 5 #define INF (1<<30)
 6 #define MAXN 55555
 7 struct Edge{
 8     int v,next;
 9 }edge[MAXN<<1];
10 int NE,head[MAXN];
11 void addEdge(int u,int v){
12     edge[NE].v=v; edge[NE].next=head[u];
13     head[u]=NE++;
14 }
15 bool vis[MAXN];
16 int size[MAXN];
17 void getSize(int u,int fa){
18     size[u]=1;
19     for(int i=head[u]; i!=-1; i=edge[i].next){
20         int v=edge[i].v;
21         if(v==fa || vis[v]) continue;
22         getSize(v,u);
23         size[u]+=size[v];
24     }
25 }
26 int cen,mini;
27 void getCentre(int u,int fa,int &tot){
28     int res=tot-size[u];
29     for(int i=head[u]; i!=-1; i=edge[i].next){
30         int v=edge[i].v;
31         if(v==fa || vis[v]) continue;
32         getCentre(v,u,tot);
33         res=max(res,size[v]);
34     }
35     if(res<mini){
36         mini=res;
37         cen=u;
38     }
39 }
40 int getCentre(int u){
41     mini=INF;
42     getSize(u,u);
43     getCentre(u,u,size[u]);
44     return cen;
45 }
46 int n,k,cnt[555],tmp[555],ans;
47 void dfs(int u,int fa,int dist){
48     if(dist>k) return;
49     ans+=cnt[k-dist];
50     ++tmp[dist];
51     for(int i=head[u]; i!=-1; i=edge[i].next){
52         int v=edge[i].v;
53         if(v==fa || vis[v]) continue;
54         dfs(v,u,dist+1);
55     }
56 }
57 void conquer(int u){
58     memset(cnt,0,sizeof(cnt));
59     cnt[0]=1;
60     for(int i=head[u]; i!=-1; i=edge[i].next){
61         int v=edge[i].v;
62         if(vis[v]) continue;
63         memset(tmp,0,sizeof(tmp));
64         dfs(v,v,1);
65         for(int i=1; i<=k; ++i) cnt[i]+=tmp[i];
66     }
67 }
68 void divide(int u){
69     u=getCentre(u);
70     vis[u]=1;
71     conquer(u);
72     for(int i=head[u]; i!=-1; i=edge[i].next){
73         int v=edge[i].v;
74         if(vis[v]) continue;
75         divide(v);
76     }
77 }
78 int main(){
79     memset(head,-1,sizeof(head));
80     int a,b;
81     scanf("%d%d",&n,&k);
82     for(int i=1; i<n; ++i){
83         scanf("%d%d",&a,&b);
84         addEdge(a,b);
85         addEdge(b,a);
86     }
87     divide(1);
88     printf("%d",ans);
89     return 0;
90 }