#include <bits/stdc++.h>
using namespace std;
#include <ext/pb_ds/tree_policy.hpp>
#include <ext/pb_ds/assoc_container.hpp>
using namespace __gnu_pbds;
template <class T> using Tree = tree<T, null_type, less_equal<T>,
rb_tree_tag, tree_order_statistics_node_update>;
#define int int64_t
#ifdef LOCAL
#include "algo/debug.h"
#endif
#define f first
#define s second
template<class T> using V = vector<T>;
using vi = V<int>;
#define all(x) begin(x), end(x)
#define rall(x) rbegin(x), rend(x)
#define pb push_back
#define lb lower_bound
#define ub upper_bound
template<class T> int lwb(V<T>& a, const T& b) { return lb(all(a),b)-begin(a); }
template<class T> int upb(V<T>& a, const T& b) { return ub(all(a),b)-begin(a); }
template<class T> bool ckmin(T& a, const T& b) { return a > b ? a=b, true : false; }
template<class T> bool ckmax(T& a, const T& b) { return a < b ? a=b, true : false; }
int n;
vi sub;
V<bool> del;
int ans;
V<V<pair<int, array<int, 2>>>> g;
V<pair<array<int, 2>, int>> paths;
void find_sub(int u, int p) {
sub[u] = 1;
for (auto [v, wi] : g[u]) if (v != p && !del[v]) {
find_sub(v, u);
sub[u] += sub[v];
}
}
int find_ctd(int u, int p, int sub_siz) {
for (auto [v, wi] : g[u]) if (v != p && !del[v]) {
if (2*sub[v] >= sub_siz) {
return find_ctd(v, u, sub_siz);
}
}
return u;
}
void find_paths(int u, int p, int d, array<int, 2> mx) {
for (auto [v, wi] : g[u]) if (v != p && !del[v]) {
ckmax(mx, wi);
paths.pb({mx, d+1});
find_paths(v, u, d+1, mx);
if (p == -1) {
mx = {0, -1};
}
}
}
int count_pairs_in_sub(int u, int p) {
int res = 0;
find_paths(u, p, 0, {0, -1});
sort(all(paths));
Tree<int> s;
for (auto [wi, d] : paths) {
int mx = wi[0];
if (mx >= d) res++;
res += s.order_of_key(mx-d+1);
s.insert(d);
}
paths.clear();
return res;
}
void process_ctd(int ctd, int p) {
// ps("processing centroid", ctd+1);
// ps("+", count_pairs_in_sub(ctd, p));
ans += count_pairs_in_sub(ctd, p);
for (auto [v, wi] : g[ctd]) if (v != p && !del[v]) {
// ps("-", count_pairs_in_sub(v, ctd));
ans -= count_pairs_in_sub(v, ctd);
}
}
void cd(int u) {
find_sub(u, -1);
int ctd = find_ctd(u, -1, sub[u]);
process_ctd(ctd, -1);
del[ctd] = true;
for (auto [v, wi] : g[ctd]) if (!del[v]) {
cd(v);
}
}
void solve() {
int k; cin >> n >> k;
g.resize(n);
del.resize(n);
sub.resize(n);
for (int i = 0; i < n-1; i++) {
int u, v, w; cin >> u >> v >> w; u--; v--; w -= k;
g[u].pb({v, {w, i}}); g[v].pb({u, {w, i}});
}
cd(0);
cout << 2*ans << '\n';
}
signed main() {
ios::sync_with_stdio(false);
cin.tie(0); cout.tie(0);
solve();
return 0;
}
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Correct |
1 ms |
348 KB |
Output is correct |
2 |
Correct |
0 ms |
348 KB |
Output is correct |
3 |
Incorrect |
1 ms |
348 KB |
Output isn't correct |
4 |
Halted |
0 ms |
0 KB |
- |
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Correct |
1 ms |
348 KB |
Output is correct |
2 |
Correct |
0 ms |
348 KB |
Output is correct |
3 |
Incorrect |
1 ms |
344 KB |
Output isn't correct |
4 |
Halted |
0 ms |
0 KB |
- |
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Correct |
1 ms |
348 KB |
Output is correct |
2 |
Correct |
0 ms |
348 KB |
Output is correct |
3 |
Incorrect |
1 ms |
348 KB |
Output isn't correct |
4 |
Halted |
0 ms |
0 KB |
- |