fork download 
  1. #include<bits/stdc++.h>
  2. using namespace std;
  3. #define int long long
  4. const int N=2e5+5;
  5. int n,k;
  6. vector<int> graph[N];
  7. int subtree[N],cnt[N]{1};
  8. int ans=0,mx_depth;
  9. bool check[N];
  10. int sz(int node,int par)
  11. {
  12.     subtree[node]=1;
  13.     for(int a:graph[node]){
  14.         if(!check[a]&&a!=par) subtree[node]+=sz(a,node);
  15.     }
  16.     return subtree[node];
  17. }
  18. int centroid(int d,int node,int par)
  19. {
  20.     for(int a:graph[node]){
  21.         if(!check[a]&&a!=par&&subtree[a]>=d) return centroid(d,a,node);
  22.     }
  23.     return node;
  24. }
  25. void num(int node,int par,bool f,int d)
  26. {
  27.     if(d>k) return;
  28.     mx_depth=max(mx_depth,d);
  29.     if(f) cnt[d]++;
  30.     else ans+=cnt[k-d];
  31.     for(int a:graph[node]){
  32.         if(!check[a]&&a!=par) num(a,node,f,d+1);
  33.     }
  34. }
  35. void cdp(int node)
  36. {
  37.     int cen=centroid(sz(node,0)/2,node,0);
  38.     check[cen]=1;
  39.     mx_depth=0;
  40.     for(int a:graph[cen]){
  41.         if(!check[a]){
  42.             num(a,cen,0,1);
  43.             num(a,cen,1,1);
  44.         }
  45.     }
  46.     fill(cnt+1,cnt+mx_depth+1,0);
  47.     for(int a:graph[cen]){
  48.         if(!check[a]) cdp(a);
  49.     }
  50. }
  51. signed main()
  52. {
  53.     ios_base::sync_with_stdio(false);
  54.     cin.tie(0);
  55.     cout.tie(0);
  56.     cin>>n>>k;
  57.     for(int a=0;a<n-1;a++){
  58.         int x,y;
  59.         cin>>x>>y;
  60.         graph[x].push_back(y);
  61.         graph[y].push_back(x);
  62.     }
  63.     cdp(1);
  64.     cout<<ans;
  65. }
  66.  
Success #stdin #stdout 0.01s 9208KB
comments (?)
stdin 
5 2
1 2
2 3
3 4
3 5
stdout 
4
https://ideone.com/5WO5wm
language:
C++14 (gcc 8.3)
created:
3 days ago
visibility:
public

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!