fork download 
  1. #include <bits/stdc++.h> 
  2.  
  3. using namespace std;  
  4.  
  5. typedef long long ll;  
  6. typedef pair<int, int> ii;  
  7.  
  8. const int INF = 1e9;  
  9. const ll LINF = 1e18;  
  10.  
  11. const int N = 2e5 + 5; 
  12. const int M = 2e5 + 5; 
  13.  
  14. // Sort lại các cạnh tăng dần theo trọng số
  15. // Khi đó xét đến cạnh thứ i thì luôn đảm bảo w(i) là lớn nhất trong số các cạnh đã được thêm
  16. // Khi thêm cạnh i thì đếm số cặp (u, v) mà có đường đi đi qua cạnh i
  17. // => DSU
  18. // O(số truy vấn * n)
  19. // Nhận xét 2: Xét 2 truy vấn i, j sao cho q(i) < q(j)
  20. // Từ truy vấn thứ i, ta có thể thêm vào một số cạnh để có được đáp án của truy vấn thứ j 
  21. // => Sort lại các truy vấn tăng dần theo q và dùng kĩ thuật Hai con trỏ
  22. // O(số truy vấn + n)
  23.  
  24. struct Edge {
  25. 	int u, v, w; 
  26. 	bool operator<(const Edge& other) const {
  27. 		return w < other.w; 
  28. 	} 
  29. }; 
  30.  
  31. struct Query {
  32. 	int q, idx;  
  33. 	bool operator<(const Query& other) const {
  34. 		return q < other.q; 
  35. 	}
  36. }; 
  37.  
  38. int n, m; 
  39. vector<Edge> edges; 
  40. vector<Query> queries;  
  41.  
  42. int p[N], sz[N]; 
  43.  
  44. void initSet() {
  45. 	for (int u = 1; u <= n; u++) {
  46. 		p[u] = u;  
  47. 		sz[u] = 1; 
  48. 	}
  49. }
  50.  
  51. int findSet(int u) {
  52. 	if (p[u] == u) return u; 
  53. 	return p[u] = findSet(p[u]); 
  54. }
  55.  
  56. ll unionSet(int u, int v) {
  57. 	u = findSet(u); 
  58. 	v = findSet(v); 
  59. 	if (u == v) return 0; 
  60. 	if (sz[u] < sz[v]) swap(u, v);  
  61. 	ll cnt = 1ll * sz[u] * sz[v]; 
  62. 	p[v] = u;  
  63. 	sz[u] += sz[v];
  64. 	return cnt;  
  65. }
  66.  
  67. ll ans[M]; 
  68.  
  69. int main() {
  70. 	ios::sync_with_stdio(false); 
  71. 	cin.tie(nullptr); 	
  72. 	cin >> n >> m; 
  73.  
  74. 	for (int i = 0; i < n - 1; i++) {
  75. 		int u, v, w; 
  76. 		cin >> u >> v >> w; 
  77. 		edges.push_back({u, v, w}); 
  78. 	}
  79.  
  80. 	for (int i = 0; i < m; i++) {
  81. 		int q; cin >> q;  
  82. 		queries.push_back({q, i});
  83. 	}
  84.  
  85. 	sort(edges.begin(), edges.end());   
  86. 	sort(queries.begin(), queries.end());  
  87.  
  88. 	initSet();  
  89.  
  90. 	ll cur_ans = 0;   
  91. 	for (int i = 0, j = 0; i < m; i++) {
  92. 		int q = queries[i].q, idx = queries[i].idx;   		
  93. 		while (j < n - 1 && edges[j].w <= q) {
  94. 			cur_ans += unionSet(edges[j].u, edges[j].v);
  95. 			j++; 
  96. 		}
  97. 		ans[idx] = cur_ans; 
  98. 	}
  99.  
  100. 	for (int i = 0; i < m; i++) cout << ans[i] << " \n"[i == m - 1]; 
  101. }
Success #stdin #stdout 0s 5288KB
comments (?)
stdin 
7 5
1 2 1
3 2 3
2 4 1
4 5 2
5 7 4
3 6 2
5 2 3 4 1
stdout 
21 7 15 21 3
https://ideone.com/GbDzIa
language:
C++14 (gcc 8.3)
created:
312 days ago
visibility:
secret

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!