#include <bits/stdc++.h>
#include<ext/pb_ds/assoc_container.hpp>
using namespace std;
using namespace __gnu_pbds;
#define int int64_t
typedef tree<pair<int, int>, null_type, less<pair<int, int>>, rb_tree_tag, tree_order_statistics_node_update> ordered_set;
const int inf = 2e9;
const int mod = 1e9 + 7;
int ans = 0;
int k;
vector<vector<pair<int, int>>> g;
vector<vector<int>> centers;
vector<bool> del;
vector<int> sz;
void find_sz(int v, int p) {
sz[v] = 1;
for (auto to_w: g[v]) {
int to = to_w.first, w = to_w.second;
if (to == p || del[to]) {
continue;
}
find_sz(to, v);
sz[v] += sz[to];
}
}
int find_center(int v, int p, int total_sz) {
for (auto to_w: g[v]) {
int to = to_w.first, w = to_w.second;
if (to == p || del[to]) {
continue;
}
if (sz[to] * 2 > total_sz) {
return find_center(to, v, total_sz);
}
}
return v;
}
void fill_list(int v, int p, int cur_max, int cur_dist, vector<pair<int, int>>& lst) {
lst.push_back({cur_max, cur_dist});
for (auto to_w: g[v]) {
int to = to_w.first, w = to_w.second;
if (to == p || del[to]) {
continue;
}
fill_list(to, v, max(cur_max, w), cur_dist + 1, lst);
}
}
void handle_center(int v) {
del[v] = true;
vector<vector<pair<int, int>>> lists; // {max_on_path, dist}
vector<pair<int, int>> common_list;
common_list.push_back({-inf, 0});
for (auto to_w: g[v]) {
int to = to_w.first, w = to_w.second;
lists.push_back(vector<pair<int, int>>(0));
if (!del[to]) {
fill_list(to, -1, w, 1, lists.back());
}
sort(lists.back().begin(), lists.back().end());
common_list.insert(common_list.end(), lists.back().begin(), lists.back().end());
}
sort(common_list.begin(), common_list.end());
auto handle_list = [&](vector<pair<int, int>>& lst) {
int counter = 0;
int cur_ans = 0;
ordered_set s;
for (auto mx_dist: lst) {
int mx = mx_dist.first, dist = mx_dist.second;
cur_ans += s.order_of_key({mx - k - dist + 1, -1});
s.insert({dist, counter++});
}
return cur_ans;
};
ans += handle_list(common_list);
for (auto lst: lists) {
ans -= handle_list(lst);
}
}
void centroid_decomposition(int v) {
find_sz(v, -1);
v = find_center(v, -1, sz[v]);
handle_center(v);
del[v] = true;
for (auto to_w: g[v]) {
int to = to_w.first, w = to_w.second;
if (!del[to]) {
centroid_decomposition(to);
}
}
}
int32_t main() {
ios_base::sync_with_stdio(0);
cin.tie(0); cout.tie(0);
int n;
cin >> n >> k;
g.resize(n);
for (int i = 0; i < n - 1; i++) {
int a, b, c;
cin >> a >> b >> c;
a--; b--;
g[a].push_back({b, c});
g[b].push_back({a, c});
}
del.resize(n);
sz.resize(n);
centers.resize(n);
centroid_decomposition(0);
cout << ans * 2 << '\n';
}
Compilation message
Main.cpp: In function 'void find_sz(int64_t, int64_t)':
Main.cpp:24:30: warning: unused variable 'w' [-Wunused-variable]
24 | int to = to_w.first, w = to_w.second;
| ^
Main.cpp: In function 'int64_t find_center(int64_t, int64_t, int64_t)':
Main.cpp:35:30: warning: unused variable 'w' [-Wunused-variable]
35 | int to = to_w.first, w = to_w.second;
| ^
Main.cpp: In function 'void centroid_decomposition(int64_t)':
Main.cpp:95:30: warning: unused variable 'w' [-Wunused-variable]
95 | int to = to_w.first, w = to_w.second;
| ^
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Correct |
0 ms |
348 KB |
Output is correct |
2 |
Correct |
0 ms |
348 KB |
Output is correct |
3 |
Correct |
1 ms |
348 KB |
Output is correct |
4 |
Correct |
2 ms |
604 KB |
Output is correct |
5 |
Correct |
2 ms |
604 KB |
Output is correct |
6 |
Correct |
2 ms |
604 KB |
Output is correct |
7 |
Correct |
2 ms |
724 KB |
Output is correct |
8 |
Correct |
3 ms |
604 KB |
Output is correct |
9 |
Correct |
1 ms |
604 KB |
Output is correct |
10 |
Correct |
1 ms |
604 KB |
Output is correct |
11 |
Correct |
2 ms |
680 KB |
Output is correct |
12 |
Correct |
2 ms |
600 KB |
Output is correct |
13 |
Correct |
3 ms |
604 KB |
Output is correct |
14 |
Correct |
2 ms |
604 KB |
Output is correct |
15 |
Correct |
2 ms |
604 KB |
Output is correct |
16 |
Correct |
2 ms |
604 KB |
Output is correct |
17 |
Correct |
2 ms |
604 KB |
Output is correct |
18 |
Correct |
2 ms |
600 KB |
Output is correct |
19 |
Correct |
2 ms |
604 KB |
Output is correct |
20 |
Correct |
2 ms |
604 KB |
Output is correct |
21 |
Correct |
2 ms |
600 KB |
Output is correct |
22 |
Correct |
4 ms |
604 KB |
Output is correct |
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Correct |
0 ms |
348 KB |
Output is correct |
2 |
Correct |
0 ms |
456 KB |
Output is correct |
3 |
Correct |
0 ms |
348 KB |
Output is correct |
4 |
Correct |
2 ms |
728 KB |
Output is correct |
5 |
Correct |
33 ms |
3176 KB |
Output is correct |
6 |
Correct |
276 ms |
14196 KB |
Output is correct |
7 |
Correct |
625 ms |
27540 KB |
Output is correct |
8 |
Correct |
714 ms |
27956 KB |
Output is correct |
9 |
Correct |
616 ms |
27704 KB |
Output is correct |
10 |
Correct |
723 ms |
28044 KB |
Output is correct |
11 |
Correct |
666 ms |
27532 KB |
Output is correct |
12 |
Correct |
736 ms |
27960 KB |
Output is correct |
13 |
Correct |
661 ms |
27696 KB |
Output is correct |
14 |
Correct |
720 ms |
28208 KB |
Output is correct |
15 |
Correct |
727 ms |
27700 KB |
Output is correct |
16 |
Correct |
656 ms |
27956 KB |
Output is correct |
17 |
Correct |
638 ms |
27968 KB |
Output is correct |
18 |
Correct |
630 ms |
27764 KB |
Output is correct |
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Correct |
0 ms |
348 KB |
Output is correct |
2 |
Correct |
0 ms |
348 KB |
Output is correct |
3 |
Correct |
1 ms |
348 KB |
Output is correct |
4 |
Correct |
2 ms |
604 KB |
Output is correct |
5 |
Correct |
2 ms |
604 KB |
Output is correct |
6 |
Correct |
2 ms |
604 KB |
Output is correct |
7 |
Correct |
2 ms |
724 KB |
Output is correct |
8 |
Correct |
3 ms |
604 KB |
Output is correct |
9 |
Correct |
1 ms |
604 KB |
Output is correct |
10 |
Correct |
1 ms |
604 KB |
Output is correct |
11 |
Correct |
2 ms |
680 KB |
Output is correct |
12 |
Correct |
2 ms |
600 KB |
Output is correct |
13 |
Correct |
3 ms |
604 KB |
Output is correct |
14 |
Correct |
2 ms |
604 KB |
Output is correct |
15 |
Correct |
2 ms |
604 KB |
Output is correct |
16 |
Correct |
2 ms |
604 KB |
Output is correct |
17 |
Correct |
2 ms |
604 KB |
Output is correct |
18 |
Correct |
2 ms |
600 KB |
Output is correct |
19 |
Correct |
2 ms |
604 KB |
Output is correct |
20 |
Correct |
2 ms |
604 KB |
Output is correct |
21 |
Correct |
2 ms |
600 KB |
Output is correct |
22 |
Correct |
4 ms |
604 KB |
Output is correct |
23 |
Correct |
0 ms |
348 KB |
Output is correct |
24 |
Correct |
0 ms |
456 KB |
Output is correct |
25 |
Correct |
0 ms |
348 KB |
Output is correct |
26 |
Correct |
2 ms |
728 KB |
Output is correct |
27 |
Correct |
33 ms |
3176 KB |
Output is correct |
28 |
Correct |
276 ms |
14196 KB |
Output is correct |
29 |
Correct |
625 ms |
27540 KB |
Output is correct |
30 |
Correct |
714 ms |
27956 KB |
Output is correct |
31 |
Correct |
616 ms |
27704 KB |
Output is correct |
32 |
Correct |
723 ms |
28044 KB |
Output is correct |
33 |
Correct |
666 ms |
27532 KB |
Output is correct |
34 |
Correct |
736 ms |
27960 KB |
Output is correct |
35 |
Correct |
661 ms |
27696 KB |
Output is correct |
36 |
Correct |
720 ms |
28208 KB |
Output is correct |
37 |
Correct |
727 ms |
27700 KB |
Output is correct |
38 |
Correct |
656 ms |
27956 KB |
Output is correct |
39 |
Correct |
638 ms |
27968 KB |
Output is correct |
40 |
Correct |
630 ms |
27764 KB |
Output is correct |
41 |
Correct |
0 ms |
348 KB |
Output is correct |
42 |
Correct |
608 ms |
27448 KB |
Output is correct |
43 |
Correct |
713 ms |
27956 KB |
Output is correct |
44 |
Correct |
642 ms |
27700 KB |
Output is correct |
45 |
Correct |
711 ms |
27884 KB |
Output is correct |
46 |
Correct |
606 ms |
27448 KB |
Output is correct |
47 |
Correct |
731 ms |
27956 KB |
Output is correct |
48 |
Correct |
632 ms |
27540 KB |
Output is correct |
49 |
Correct |
719 ms |
27956 KB |
Output is correct |
50 |
Correct |
733 ms |
27700 KB |
Output is correct |
51 |
Correct |
644 ms |
27960 KB |
Output is correct |
52 |
Correct |
164 ms |
22160 KB |
Output is correct |
53 |
Correct |
182 ms |
22340 KB |
Output is correct |
54 |
Correct |
188 ms |
22060 KB |
Output is correct |
55 |
Correct |
180 ms |
22380 KB |
Output is correct |
56 |
Correct |
200 ms |
22272 KB |
Output is correct |
57 |
Correct |
434 ms |
21684 KB |
Output is correct |
58 |
Correct |
435 ms |
22012 KB |
Output is correct |
59 |
Correct |
563 ms |
22108 KB |
Output is correct |
60 |
Correct |
493 ms |
22024 KB |
Output is correct |
61 |
Correct |
483 ms |
22060 KB |
Output is correct |
62 |
Correct |
342 ms |
21880 KB |
Output is correct |
63 |
Correct |
419 ms |
22140 KB |
Output is correct |
64 |
Correct |
417 ms |
22176 KB |
Output is correct |
65 |
Correct |
13 ms |
1372 KB |
Output is correct |
66 |
Correct |
0 ms |
348 KB |
Output is correct |