fork download
copy 
  1. // base off ecnerwala's solution
  2. #ifndef LOCAL
  3. #define NDEBUG
  4. #endif
  5. #include <bits/stdc++.h>
  6. #include <bits/extc++.h>
  7. using namespace std;
  8.  
  9. namespace __gnu_pbds {
  10. template <typename K, typename V=null_type>
  11. using order_statistic_tree = tree<K, V, less<K>, rb_tree_tag, tree_order_statistics_node_update>;
  12. }
  13. using __gnu_pbds::order_statistic_tree;
  14.  
  15. using i64 = int64_t;
  16.  
  17. const int MAXN = 1.6e5;
  18. int N, M, K;
  19. vector<int> adj[MAXN];
  20.  
  21. vector<int> preorder;
  22. int sz[MAXN];
  23. int depth[MAXN];
  24. int par[MAXN];
  25.  
  26. int heavyRoot[MAXN];
  27. int st[MAXN];
  28. int en[MAXN];
  29.  
  30. void dfs_sz(int cur, int prv) {
  31. 	if (prv != -1) {
  32. 		adj[cur].erase(find(adj[cur].begin(), adj[cur].end(), prv));
  33. 	}
  34. 	par[cur] = prv;
  35. 	depth[cur] = (prv == -1 ? 0 : depth[prv]+1);
  36. 	sz[cur] = 1;
  37.  
  38. 	for (int nxt : adj[cur]) {
  39. 		dfs_sz(nxt, cur);
  40. 		sz[cur] += sz[nxt];
  41. 	}
  42.  
  43. 	nth_element(adj[cur].begin(), adj[cur].begin(), adj[cur].end(), [&](int i, int j) { return sz[i] > sz[j]; });
  44. }
  45.  
  46. void dfs_hld(int cur) {
  47. 	st[cur] = int(preorder.size());
  48. 	preorder.push_back(cur);
  49.  
  50. 	heavyRoot[cur] = (par[cur] != -1 && st[cur] == st[par[cur]]+1) ? heavyRoot[par[cur]] : cur;
  51.  
  52. 	for (int nxt : adj[cur]) {
  53. 		dfs_hld(nxt);
  54. 	}
  55.  
  56. 	en[cur] = int(preorder.size());
  57. }
  58.  
  59. int lca(int a, int b) {
  60. 	while (heavyRoot[a] != heavyRoot[b]) {
  61. 		// will st[a] > st[b] work? 
  62. 		if (st[a] > st[b]) {
  63. 			swap(a, b);
  64. 		}
  65. 		b = par[heavyRoot[b]];
  66. 	}
  67. 	return st[a] < st[b] ? a : b;
  68. }
  69.  
  70. int get_anc(int cur, int k) {
  71. 	assert(k >= 0);
  72. 	int d = depth[cur] - k;
  73. 	while (d < depth[heavyRoot[cur]]) {
  74. 		cur = par[heavyRoot[cur]];
  75. 	}
  76. 	return preorder[st[cur] - (depth[cur] - d)];
  77. }
  78.  
  79. using evt_t = pair<pair<int, int>, int>;
  80.  
  81. vector<evt_t> evts[MAXN];
  82. i64 ans = 0;
  83.  
  84. using tree_t = order_statistic_tree<evt_t>;
  85.  
  86. tree_t dfs(int cur) {
  87. 	tree_t res;
  88.  
  89. 	int curPar = depth[cur] >= K ? get_anc(cur, K-1) : -1;
  90.  
  91. 	auto joinEvt = [&](pair<int, int> it) -> bool {
  92. 		auto getSubtree = [&](int n) -> int {
  93. 			return int(res.order_of_key({{en[n], -1}, -1})) - int(res.order_of_key({{st[n], -1}, -1}));
  94. 		};
  95. 		int stDest = it.first, c = it.second;
  96. 		assert(!(st[cur] <= stDest && stDest < en[cur]));
  97. 		int dest = preorder[stDest];
  98.  
  99. 		int totDist = depth[cur] + depth[dest] - 2 * depth[c];
  100. 		if (totDist >= K) {
  101. 			if (depth[cur] - depth[c] >= K) {
  102. 				assert(curPar != -1);
  103. 				ans += int(res.size()) - getSubtree(curPar);
  104. 			} else {
  105. 				ans += getSubtree(get_anc(dest, totDist - K));
  106. 			}
  107. 			if (stDest >= en[cur]) {
  108. 				// event of the second kind
  109. 				assert(par[dest] == c);
  110. 				ans -= getSubtree(dest);
  111. 			}
  112. 			return true;
  113. 		} else {
  114. 			return false;
  115. 		}
  116. 	};
  117.  
  118. 	for (auto it : evts[cur]) {
  119. 		if (joinEvt(it.first)) {
  120. 			res.insert(it);
  121. 		}
  122. 	}
  123. 	for (int nxt : adj[cur]) {
  124. 		auto o = dfs(nxt);
  125. 		for (auto lo = o.lower_bound({{st[cur], -1}, -1}); lo != o.end() && lo->first.first < en[cur]; ) {
  126. 			lo = o.erase(lo);
  127. 		}
  128. 		if (o.size() > res.size()) res.swap(o);
  129. 		for (auto it = o.begin(); it != o.end(); ) {
  130. 			if (joinEvt(it->first)) {
  131. 				++it;
  132. 			} else {
  133. 				it = o.erase(it);
  134. 			}
  135. 		}
  136. 		for (auto evt : o) {
  137. 			res.insert(evt);
  138. 		}
  139. 	}
  140.  
  141. 	return res;
  142. }
  143.  
  144. int main() {
  145. 	ios::sync_with_stdio(0), cin.tie(0);
  146. 	cin >> N >> M >> K;
  147. 	for (int e = 0; e < N-1; e++) {
  148. 		int u, v; cin >> u >> v; u--, v--;
  149. 		adj[u].push_back(v);
  150. 		adj[v].push_back(u);
  151. 	}
  152.  
  153. 	dfs_sz(0, -1);
  154. 	preorder.reserve(N);
  155. 	dfs_hld(0);
  156. 	assert(int(preorder.size()) == N);
  157.  
  158. 	int evtIdx = 0;
  159. 	for (int i = 0; i < M; i++) {
  160. 		int s, t; cin >> s >> t; s--, t--;
  161. 		if (st[s] > st[t]) swap(s, t);
  162. 		int c = lca(s, t);
  163. 		evts[t].push_back({{st[s], c}, evtIdx++});
  164. 		if (s != c) {
  165. 			assert(depth[c] < depth[t]);
  166. 			evts[s].push_back({{st[get_anc(t, depth[t] - depth[c] - 1)], c}, evtIdx++});
  167. 		}
  168. 	}
  169.  
  170. 	dfs(0);
  171. 	cout << ans << '\n';
  172.  
  173. 	return 0;
  174. }
  175.  
Success #stdin #stdout 0s 11248KB
comments (?)
stdin
copy 
13 8 3
7 6
9 11
5 6
11 3
9 7
2 12
4 3
1 2
5 8
6 13
5 10
3 1
10 4
10 11
8 11
4 9
2 5
3 5
7 3
8 10
stdout
copy 
14
https://ideone.com/l8PvN7
language:
C++14 (gcc 8.3)
created:
5 years 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!