[codeforces161D]Distance in Tree(点分治)

  1 #pragma comment(linker,"/STACK:102400000,102400000")
  2 #include<cstdio>
  3 #include<cstring>
  4 #include<algorithm>
  5 #include<cstdlib>
  6 #include<iostream>
  7 #include<cmath>
  8 #include<map>
  9 #define maxn 100040
 10 #define maxm 1000500
 11 using namespace std;
 12 typedef long long ll;
 13 const ll mod=1000003;
 14 const ll inf=1ll<<60;
 15 ll n,k,ans,size,s[maxn],f[maxn],path[maxn],cr;
 16 ll head[maxn],cnt,root;
 17 bool vis[maxn];
 18 struct edge{
 19     ll to,nxt;
 20 }e[maxn<<1];
 21 map<int,int>mp;
 22 void add_edge(ll u,ll v){
 23     e[cnt].to=v;
 24     e[cnt].nxt=head[u];
 25     head[u]=cnt++;
 26 }
 27 
 28 inline ll read(){
 29     char k=0;char ls;ls=getchar();for(;ls<‘0‘||ls>‘9‘;k=ls,ls=getchar());
 30     ll x=0;for(;ls>=‘0‘&&ls<=‘9‘;ls=getchar())x=(x<<3)+(x<<1)+ls-‘0‘;
 31     if(k==‘-‘)x=0-x;return x;
 32 }
 33 
 34 void get_root(ll u,ll fa){//get_root会用到size
 35     s[u]=1;f[u]=0;//f是dp数组
 36     for(ll i=head[u];i!=-1;i=e[i].nxt){
 37         ll v=e[i].to;
 38         if(v==fa||vis[v]) continue;
 39         get_root(v,u);
 40         s[u]+=s[v];
 41         f[u]=max(f[u],s[v]);
 42     }
 43     f[u]=max(f[u],size-s[u]);
 44     root=f[root]>f[u]?u:root;
 45 }
 46 
 47 void get_path_size(ll u,ll fa,ll dis){
 48     if(dis+1<=k){
 49         path[cr]=dis+1;
 50         cr++;
 51     }
 52     s[u]=1;
 53     for(ll i=head[u];i!=-1;i=e[i].nxt){
 54         ll v=e[i].to;
 55         if(v==fa||vis[v]) continue;
 56         get_path_size(v,u,dis+1);
 57         s[u]+=s[v];
 58     }
 59 }
 60 
 61 void work(ll u,ll fa){
 62     vis[u]=true;
 63     mp.clear();
 64     mp[0]=1;
 65     for(ll i=head[u];i!=-1;i=e[i].nxt){
 66         ll v=e[i].to;
 67         if(v==fa||vis[v]) continue;
 68         cr=0;
 69         get_path_size(v,u,0);
 70         for(ll j=0;j<cr;j++){
 71             ans+=mp[k-path[j]];
 72         }
 73         for(int j=0;j<cr;j++){
 74             mp[path[j]]++;
 75         }
 76     }
 77     for(ll i=head[u];i!=-1;i=e[i].nxt){
 78         ll v=e[i].to;
 79         if(vis[v]||v==fa) continue;
 80         size=s[v],root=0;
 81         get_root(v,u);
 82         work(root,u);
 83     }
 84 }
 85 
 86 void init(){
 87     memset(vis,0,sizeof vis);
 88     memset(head,-1,sizeof head);
 89     ans=cnt=0;
 90 }
 91 
 92 int main(){
 93     ll a,b;
 94     f[0]=inf;
 95     while(scanf("%I64d%I64d",&n,&k)!=EOF){
 96         init();
 97         for(int i=0;i<n-1;i++){
 98             a=read(),b=read();
 99             add_edge(a,b);
100             add_edge(b,a);
101         }
102         size=n,root=0;
103         get_root(1,-1);
104         work(root,-1);
105         printf("%d\n",ans);
106     }
107     return 0;
108 }