fork download
copy 
  1. // author : Nguyễn Trọng Nguyễn - ITK22 NBK 
  2. #include <bits/stdc++.h>
  3.  
  4. #define ll long long
  5. #define ii pair <int, int>
  6. #define fi first
  7. #define sc second
  8.  
  9. using namespace std;
  10.  
  11. const int maxn = (int)2e5;
  12. const int BS = 320;
  13. const int MOD = (int)1e9 + 7;
  14.  
  15. int n, k, mx;
  16. vector <int> adj[maxn + 5];
  17. int cnt[maxn + 5], sz[maxn + 5];
  18. bool vist[maxn + 5];
  19. ll ans = 0;
  20.  
  21. int subtree_size (int u, int p = 0) {
  22. 	sz[u] = 1;
  23. 	for (auto v : adj[u]) {
  24. 		if (v == p or vist[v]) continue;
  25. 		sz[u] += subtree_size(v, u);
  26. 	}
  27. 	return sz[u];
  28. }
  29.  
  30. int find_centroid (int tree_size, int u, int p = 0) {
  31. 	for (auto v : adj[u]) {
  32. 		if (vist[v] or v == p) continue;
  33. 		if (sz[v] > tree_size / 2) return find_centroid(tree_size, v, u);
  34. 	}
  35. 	return u;
  36. }
  37.  
  38. void get (bool type, int u, int p, int dist = 1) {
  39. 	if (dist > k) return ;
  40. 	mx = max(mx, dist);
  41.  
  42. 	if (!type) cnt[dist]++;
  43. 	else ans += cnt[k - dist];
  44.  
  45. 	for (auto v : adj[u]) {
  46. 		if (vist[v] or v == p) continue;
  47. 		get(type, v, u, dist + 1);
  48. 	}
  49. }
  50.  
  51. void centroid_decompose (int u) {
  52. 	int centroid = find_centroid(subtree_size(u), u);
  53. 	vist[centroid] = true;
  54. 	mx = 0;
  55.  
  56. 	for (auto v : adj[centroid]) {
  57. 		if (vist[v]) continue;
  58. 		get(true, v, centroid);
  59. 		get(false, v, centroid);
  60. 	}
  61.  
  62. 	fill(cnt + 1, cnt + 1 + mx, 0);
  63. 	for (auto v : adj[centroid]) {
  64. 		if (vist[v]) continue;
  65. 		centroid_decompose(v);
  66. 	}
  67. }
  68.  
  69. signed main (void) {
  70. 	cin.tie(0)->sync_with_stdio(false);
  71.  
  72. 	#ifndef ONLINE_JUDGE
  73. 	freopen("test.inp","r",stdin);
  74. 	freopen("test.out","w",stdout);
  75. 	#endif
  76.  
  77. 	cin >> n >> k;
  78. 	for (int i = 1; i < n; i++) {
  79. 		int u, v; cin >> u >> v;
  80. 		adj[u].push_back(v);
  81. 		adj[v].push_back(u);
  82. 	}	
  83.  
  84. 	cnt[0] = 1;
  85. 	centroid_decompose(1);
  86. 	cout << ans;
  87.  
  88. 	return 0;
  89. }
Success #stdin #stdout 0.01s 9424KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
Standard output is empty
https://ideone.com/aUWD79
language:
C++14 (gcc 8.3)
created:
1 day 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!