#ifdef LOCAL
#include "local.h"
#else
#include <bits/stdc++.h>
#define debug(...)
#define DB(...)
#pragma GCC optimize("O3,unroll-loops")
#pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")
#endif
using namespace std;
static const bool __initialization = []() {
cin.tie(nullptr)->sync_with_stdio(false);
#define TASK "JANJETINA"
if (fopen(TASK".inp", "r")) {
(void)(!freopen(TASK".inp", "r", stdin));
(void)(!freopen(TASK".out", "w", stdout)); }
return true;
}();
using ll = long long;
using ld = long double;
#define fi first
#define se second
#define all(x) x.begin(), x.end()
#define rall(x) x.rbegin(), x.rend()
#define For(i, l, r) for (int i = (l); i <= (r); ++i)
#define Ford(i, r, l) for (int i = (r); i >= (l); --i)
#define Rep(i, n) For (i, 0, (n) - 1)
#define Repd(i, n) Ford (i, (n) - 1, 0)
#define Fe(...) for (auto __VA_ARGS__)
template <class C> inline int isz(const C& c) { return static_cast<int>(c.size()); }
template <class T> inline bool chmin(T& a, const T& b) { return (a > b) ? a = b, true : false; }
template <class T> inline bool chmax(T& a, const T& b) { return (a < b) ? a = b, true : false; }
constexpr ld eps = 1e-9;
constexpr int inf = 1e9;
constexpr ll linf = 1e18;
// =============================================================================
constexpr int maxn = 1e5 + 5;
struct Fen {
int fen[maxn];
Fen() {
memset(fen, 0, sizeof(fen));
}
void update(int i, int x) {
for (; i < maxn; i += (i & -i)) {
fen[i] += x;
}
}
int get(int i) {
int ret = 0;
for (; i > 0; i -= (i & -i)) {
ret += fen[i];
}
return ret;
}
};
struct Link { int v, w; };
int n, k;
vector<Link> g[maxn];
int sz[maxn];
bool rem[maxn];
int getSz(int u, int p = -1) {
sz[u] = 1;
Fe (&[v, _] : g[u]) {
if (!rem[v] && v != p) {
sz[u] += getSz(v, u);
}
}
return sz[u];
}
int getCen(int target, int u, int p = -1) {
Fe (&[v, _] : g[u]) {
if (rem[v] || v == p) continue;
if (sz[v] * 2 > target) {
return getCen(target, v, u);
}
}
return u;
}
vector<pair<int, int>> other;
Fen fen;
ll ans = 0;
void findSub(int u, int p, int maxw, int dep) {
other.emplace_back(maxw, dep);
Fe (&[v, w] : g[u]) {
if (!rem[v] && v != p) {
findSub(v, u, max(maxw, w), dep + 1);
}
}
}
ll calc() {
ll ret = 0;
sort(all(other));
// max(w1, w2) - (l1 + l2) >= k
// max(w1, w2) = w2
// <=> w2 - l1 - l2 >= k
// <=> w2 - l2 >= k + l1
Fe (&[maxw, dep] : other) {
ret += fen.get(maxw - dep);
fen.update(dep + k, 1);
}
Fe (&[_, dep] : other) {
fen.update(dep + k, -1);
}
return ret;
}
void solve(int u) {
int cen = getCen(getSz(u), u);
rem[cen] = true;
other.clear();
findSub(cen, -1, 0, 0);
ans += calc();
// only consider paths on different subtrees
// => eliminate paths on the same subtree
Fe (&[v, w] : g[cen]) {
if (!rem[v]) {
other.clear();
findSub(v, u, w, 1);
ans -= calc();
}
}
Fe (&[v, _] : g[cen]) {
if (!rem[v]) {
solve(v);
}
}
}
signed main() {
cin >> n >> k;
Rep (_, n - 1) {
int u, v, w; cin >> u >> v >> w;
g[u].push_back({v, w});
g[v].push_back({u, w});
}
solve(1);
cout << ans * 2 << '\n';
}
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
3 ms |
3020 KB |
Output is correct |
| 2 |
Correct |
1 ms |
3020 KB |
Output is correct |
| 3 |
Incorrect |
2 ms |
3020 KB |
Output isn't correct |
| 4 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
2 ms |
3020 KB |
Output is correct |
| 2 |
Correct |
2 ms |
3020 KB |
Output is correct |
| 3 |
Incorrect |
2 ms |
3020 KB |
Output isn't correct |
| 4 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
3 ms |
3020 KB |
Output is correct |
| 2 |
Correct |
1 ms |
3020 KB |
Output is correct |
| 3 |
Incorrect |
2 ms |
3020 KB |
Output isn't correct |
| 4 |
Halted |
0 ms |
0 KB |
- |
