fork download
copy 
  1. #include<bits/stdc++.h>
  2. using namespace std;
  3. #define int long long
  4. #define fi first
  5. #define se second
  6. #define siz(x) (int)(x.size())
  7. #define all(x) x.begin(), x.end()
  8. #define debug_arr(x,len) for(int _=1; _<=len; _++) cout<<x[_]<<" "; cout<<'\n';
  9. #define debug(x) cout<<'\n'<<#x<<": "<<x<<'\n';
  10. const int maxN  = 2e5+5;
  11.  
  12. int n, k, sz[maxN], cnt[maxN], mx_depth = 0, ans = 0, tree[maxN];
  13. bool del[maxN];
  14. vector<int>adj[maxN];
  15.  
  16. void dfs(int u, int v)
  17. {
  18. 	sz[u] = 1;
  19. 	for(auto i: adj[u])
  20. 	{
  21. 		if(i == v || del[i]) continue;
  22. 		dfs(i, u);
  23. 		sz[u] += sz[i];
  24. 	}
  25. }
  26.  
  27. int find_centroid(int u, int v, int total)
  28. {
  29. 	for(auto i: adj[u])
  30. 	{
  31. 		if(i == v || del[i]) continue;
  32. 		if(sz[i] > total/2) return find_centroid(i, u, total);
  33. 	}
  34. 	return u;
  35. }
  36.  
  37. void update(int u, int v, int layer, int type)
  38. {
  39. 	if(layer > k) return;
  40. 	mx_depth = max(mx_depth, layer);
  41. 	if(type == 1) cnt[layer]++;
  42. 	else ans += cnt[k-layer];
  43. 	for(auto i: adj[u])
  44. 	{
  45. 		if(i == v || del[i]) continue;
  46. 		update(i, u, layer+1, type);
  47. 	}
  48. }
  49.  
  50. void cal(int u)
  51. {
  52. 	dfs(u, 0);
  53. 	int root = find_centroid(u, 0, sz[u]);
  54. 	mx_depth = 0; cnt[0] = 1; del[root] = 1;
  55. 	for(auto i: adj[root])
  56. 	{
  57. 		if(del[i]) continue;
  58. 		update(i, root, 1, 0);
  59. 		update(i, root, 1, 1);
  60. 	}
  61. 	fill(cnt+1, cnt+mx_depth+1, 0);
  62. 	for(auto i: adj[root])
  63. 	{
  64. 		if(del[i]) continue;
  65. 		cal(i);
  66. 	}
  67. }
  68. int32_t main()
  69. {
  70. 	ios_base::sync_with_stdio(0); cin.tie(0);
  71. 	cin>>n>>k;
  72. 	for(int i=1; i<n; i+=1)
  73. 	{
  74. 		int x,y; cin>>x>>y;
  75. 		adj[x].push_back(y);
  76. 		adj[y].push_back(x);
  77. 	}
  78. 	dfs(1, 0);
  79. 	cal(1);
  80. 	cout<<ans<<'\n';
  81. }
Success #stdin #stdout 0.01s 12692KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
0
https://ideone.com/yiVrAg
language:
C++14 (gcc 8.3)
created:
2 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!