#include <bits/stdc++.h>
using std::vector;
using std::array;
using std::pair;
using std::tuple;
template <class F> struct RecLambda : private F {
explicit RecLambda(F&& f) : F(std::forward<F>(f)) {}
template <class... Args> decltype(auto) operator()(Args&&... args) const {
return F::operator()(*this, std::forward<Args>(args)...);
}
};
using ll = long long;
struct Fenwick {
int size;
vector<int> data;
explicit Fenwick(const int n) : size(n), data(n + 1) {}
void add(int i, const int x) {
i += 1;
while (i <= size) {
data[i] += x;
i += i & -i;
}
}
int pref(int i) const {
i += 1;
int ret = 0;
while (i > 0) {
ret += data[i];
i -= i & -i;
}
return ret;
}
};
int main() {
std::ios_base::sync_with_stdio(false);
std::cin.tie(nullptr);
int N, K;
std::cin >> N >> K;
vector<vector<pair<int, int>>> graph(N);
for (int i = 1; i < N; ++i) {
int x, y, w;
std::cin >> x >> y >> w;
x -= 1, y -= 1;
graph[x].emplace_back(y, w);
graph[y].emplace_back(x, w);
}
Fenwick fen(N);
const auto count = [&](vector<pair<int, int>>& yx) {
std::sort(yx.begin(), yx.end());
ll ret = 0;
for (const auto& [y, x] : yx) {
ret += fen.pref(std::min(N - 1, y - x - K));
fen.add(x, 1);
}
for (const auto& [y, x] : yx) {
fen.add(x, -1);
}
return ret;
};
vector<char> dead(N);
vector<int> subtree(N);
ll ans = 0;
RecLambda([&](auto&& dfs, const int root) -> void {
RecLambda([&](auto&& dfs, const int u, const int p) -> void {
subtree[u] = 1;
for (const auto& [v, w] : graph[u]) {
if (!dead[v] and v != p) {
dfs(v, u);
subtree[u] += subtree[v];
}
}
})(root, -1);
const int cent = RecLambda([&](auto&& dfs, const int u, const int p) -> int {
for (const auto& [v, w] : graph[u]) {
if (!dead[v] and v != p and subtree[v] * 2 > subtree[root]) {
return dfs(v, u);
}
}
return u;
})(root, -1);
vector<pair<int, int>> list;
vector<pair<int, int>> all;
const auto enumerate = RecLambda([&](auto&& dfs, const int u, const int p, const int x, const int y) -> void {
if (y - x >= K) {
ans += 1;
}
list.emplace_back(y, x);
all.emplace_back(y, x);
for (const auto& [v, w] : graph[u]) {
if (!dead[v] and v != p) {
return dfs(v, u, x + 1, std::max(y, w));
}
}
});
for (const auto& [u, w] : graph[cent]) {
if (!dead[u]) {
list.clear();
enumerate(u, cent, 1, w);
ans -= count(list);
}
}
ans += count(all);
dead[cent] = true;
for (const auto& [u, w] : graph[cent]) {
if (!dead[u]) {
dfs(u);
}
}
})(0);
std::cout << ans * 2 << '\n';
return 0;
}
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
1 ms |
312 KB |
Output is correct |
2 |
Correct |
0 ms |
204 KB |
Output is correct |
3 |
Correct |
1 ms |
204 KB |
Output is correct |
4 |
Correct |
2 ms |
332 KB |
Output is correct |
5 |
Correct |
3 ms |
332 KB |
Output is correct |
6 |
Correct |
2 ms |
332 KB |
Output is correct |
7 |
Correct |
3 ms |
332 KB |
Output is correct |
8 |
Correct |
2 ms |
332 KB |
Output is correct |
9 |
Correct |
1 ms |
320 KB |
Output is correct |
10 |
Incorrect |
1 ms |
320 KB |
Output isn't correct |
11 |
Halted |
0 ms |
0 KB |
- |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
1 ms |
204 KB |
Output is correct |
2 |
Correct |
0 ms |
204 KB |
Output is correct |
3 |
Correct |
1 ms |
204 KB |
Output is correct |
4 |
Correct |
3 ms |
448 KB |
Output is correct |
5 |
Correct |
25 ms |
1720 KB |
Output is correct |
6 |
Correct |
162 ms |
7476 KB |
Output is correct |
7 |
Correct |
336 ms |
14348 KB |
Output is correct |
8 |
Correct |
426 ms |
14844 KB |
Output is correct |
9 |
Correct |
335 ms |
14440 KB |
Output is correct |
10 |
Correct |
425 ms |
14740 KB |
Output is correct |
11 |
Correct |
312 ms |
14356 KB |
Output is correct |
12 |
Correct |
374 ms |
14888 KB |
Output is correct |
13 |
Correct |
347 ms |
14416 KB |
Output is correct |
14 |
Correct |
392 ms |
14744 KB |
Output is correct |
15 |
Correct |
406 ms |
14588 KB |
Output is correct |
16 |
Correct |
374 ms |
14732 KB |
Output is correct |
17 |
Correct |
372 ms |
14736 KB |
Output is correct |
18 |
Correct |
383 ms |
14724 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
1 ms |
312 KB |
Output is correct |
2 |
Correct |
0 ms |
204 KB |
Output is correct |
3 |
Correct |
1 ms |
204 KB |
Output is correct |
4 |
Correct |
2 ms |
332 KB |
Output is correct |
5 |
Correct |
3 ms |
332 KB |
Output is correct |
6 |
Correct |
2 ms |
332 KB |
Output is correct |
7 |
Correct |
3 ms |
332 KB |
Output is correct |
8 |
Correct |
2 ms |
332 KB |
Output is correct |
9 |
Correct |
1 ms |
320 KB |
Output is correct |
10 |
Incorrect |
1 ms |
320 KB |
Output isn't correct |
11 |
Halted |
0 ms |
0 KB |
- |