#include <bits/stdc++.h>
#define fastio ios_base::sync_with_stdio(0), cin.tie(0)
#define endl '\n'
#define fi first
#define se second
#define sz(v) (int)(v).size()
#define all(v) (v).begin(), (v).end()
#define compress(v) sort(all(v)), (v).erase(unique(all((v))), v.end())
#define pb push_back
#define eb emplace_back
using namespace std;
using ll = long long;
using pii = pair<int, int>;
using pll = pair<ll, ll>;
using ld = long double;
template <typename T>
inline void Mx(T &x, T y) { x < y && (x = y); }
template <typename T>
inline void Mn(T &x, T y) { x > y && (x = y); }
template <typename T>
struct Fenwick {
int n;
vector<T> tree;
Fenwick(int _n) : n(_n) {
tree.resize(n);
}
void update(int x, T v) {
for (int i = x; i < n; i += i & -i)
tree[i] += v;
}
T query(int x) {
T ret = 0;
for (int i = x; i > 0; i -= i & -i)
ret += tree[i];
return ret;
}
T query(int l, int r) {
return query(r) - query(l - 1);
}
};
const int MX = 1e5 + 5;
int N, K;
ll ans;
vector<pii> adj[MX];
int t_siz[MX];
bool fin[MX];
Fenwick<int> tree_cnt(MX);
int qn;
array<int, 3> que[MX];
int getCent(int cur, int pv, int t_size) {
for (auto &nxt : adj[cur]) {
if (nxt.first == pv || fin[nxt.first]) {
continue;
}
if (t_size < 2 * t_siz[nxt.first]) {
return getCent(nxt.first, cur, t_size);
}
}
return cur;
}
void getSz(int cur, int pv) {
t_siz[cur] = 1;
for (auto &nxt : adj[cur]) {
if (nxt.first == pv || fin[nxt.first]) {
continue;
}
getSz(nxt.first, cur);
t_siz[cur] += t_siz[nxt.first];
}
}
int dfs(int cur, int pv, int type, int mx, int d) {
int ret = 1;
que[qn++] = { mx, d, type };
tree_cnt.update(d, 1);
for (auto &nxt : adj[cur]) {
if (nxt.first == pv || fin[nxt.first]) {
continue;
}
ret += dfs(nxt.first, cur, type, max(mx, nxt.second), d + 1);
}
return ret;
}
void solve(int node) {
getSz(node, -1);
int cent = getCent(node, -1, t_siz[node]);
fin[cent] = true;
qn = 0;
vector<int> S;
vector<Fenwick<int>> rem;
for (auto &nxt : adj[cent]) {
if (fin[nxt.first]) {
continue;
}
S.pb(dfs(nxt.first, -1, sz(S), nxt.second, 1));
int s = S.back();
rem.pb(Fenwick<int>(s + 1));
int z = sz(rem) - 1;
for (int i = qn - s; i < qn; i++) {
rem[z].update(que[i][1], 1);
}
}
sort(que, que + qn);
for (int i = qn - 1; i >= 0; i--) {
auto &[mx, d, type] = que[i];
int T = mx - (d + K);
if (T >= 0) {
ans += 1;
if (T > 0) {
ans += tree_cnt.query(1, T);
ans -= rem[type].query(1, T);
}
}
tree_cnt.update(d, -1);
rem[type].query(d, -1);
}
for (auto &nxt : adj[cent]) {
if (fin[nxt.first]) {
continue;
}
solve(nxt.first);
}
}
int main() {
fastio;
cin >> N >> K;
for (int i = 0; i < N - 1; i++) {
int u, v, w;
cin >> u >> v >> w;
adj[u].eb(v, w);
adj[v].eb(u, w);
}
solve(1);
cout << ans * 2 << endl;
return 0;
}
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
2 ms |
3028 KB |
Output is correct |
| 2 |
Correct |
2 ms |
3028 KB |
Output is correct |
| 3 |
Incorrect |
2 ms |
3028 KB |
Output isn't correct |
| 4 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
2 ms |
3064 KB |
Output is correct |
| 2 |
Correct |
2 ms |
3084 KB |
Output is correct |
| 3 |
Incorrect |
3 ms |
3172 KB |
Output isn't correct |
| 4 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
2 ms |
3028 KB |
Output is correct |
| 2 |
Correct |
2 ms |
3028 KB |
Output is correct |
| 3 |
Incorrect |
2 ms |
3028 KB |
Output isn't correct |
| 4 |
Halted |
0 ms |
0 KB |
- |
