#include <algorithm>
#include <iostream>
#include <numeric>
#include <cassert>
#include <vector>
#define int long long
typedef long long llong;
const int MAXN = 100000 + 10;
const llong INF = 1e18;
const int INTINF = 1e9;
int n, k;
struct SegmentTree
{
struct Node
{
llong sum;
llong max;
llong min;
int idxMin;
int idx;
Node()
{
max = -1;
min = INF;
idxMin = 0;
idx = 0;
sum = 0;
}
};
Node tree[4*MAXN];
Node combine(Node left, Node right)
{
Node res;
res.sum = left.sum + right.sum;
if (left.max > right.max)
{
res.max = left.max;
res.idx = left.idx;
} else
{
res.max = right.max;
res.idx = right.idx;
}
if (left.min < right.min)
{
res.min = left.min;
res.idxMin = left.idxMin;
} else
{
res.min = right.min;
res.idxMin = right.idxMin;
}
return res;
}
void build(int l, int r, int node, bool inactive)
{
if (l == r)
{
tree[node].max = (inactive ? -1 : 0);
tree[node].sum = 0;
tree[node].min = INF;
tree[node].idx = l;
tree[node].idxMin = l;
return;
}
int mid = l + r >> 1;
build(l, mid, 2*node, inactive);
build(mid + 1, r, 2*node + 1, inactive);
tree[node] = combine(tree[2*node], tree[2*node + 1]);
}
void updateSet(int l, int r, int node, int queryPos, llong queryValue)
{
if (l == r)
{
if (queryValue == -1)
{
tree[node].max = -1;
tree[node].sum = 0;
tree[node].min = INF;
} else
{
tree[node].max = queryValue;
tree[node].min = queryValue;
tree[node].sum = queryValue;
}
return;
}
int mid = l + r >> 1;
if (queryPos <= mid) updateSet(l, mid, 2*node, queryPos, queryValue);
else updateSet(mid + 1, r, 2*node + 1, queryPos, queryValue);
tree[node] = combine(tree[2*node], tree[2*node + 1]);
}
void updateAdd(int l, int r, int node, int queryPos, llong queryValue)
{
if (l == r)
{
tree[node].max += queryValue;
tree[node].min += queryValue;
tree[node].sum += queryValue;
return;
}
int mid = l + r >> 1;
if (queryPos <= mid) updateAdd(l, mid, 2*node, queryPos, queryValue);
else updateAdd(mid + 1, r, 2*node + 1, queryPos, queryValue);
tree[node] = combine(tree[2*node], tree[2*node + 1]);
}
Node query(int l, int r, int node, int queryL, int queryR)
{
if (queryL <= l && r <= queryR)
{
return tree[node];
}
int mid = l + r >> 1;
Node res;
if (queryL <= mid) res = combine(res, query(l, mid, 2*node, queryL, queryR));
if (mid + 1 <= queryR) res = combine(res, query(mid + 1, r, 2*node + 1, queryL, queryR));
return res;
}
void build(bool inactive)
{
build(1, n, 1, inactive);
}
void updateSet(int pos, llong value)
{
updateSet(1, n, 1, pos, value);
}
void updateAdd(int pos, llong value)
{
updateAdd(1, n, 1, pos, value);
}
Node query(int l, int r)
{
return query(1, n, 1, l, r);
}
};
struct Fenwick
{
int tree[MAXN];
void update(int pos, int value)
{
for (int idx = pos ; idx <= n ; idx += idx & (-idx))
{
tree[idx] += value;
}
}
int query(int pos)
{
int res = 0;
for (int idx = pos ; idx > 0 ; idx -= idx & (-idx))
{
res += tree[idx];
}
return res;
}
int query(int l, int r)
{
return query(r) - query(l - 1);
}
};
int timer;
int in[MAXN];
int out[MAXN];
SegmentTree treeActive;
SegmentTree treeNotActive;
std::vector <std::pair <int,int>> g[MAXN];
Fenwick fenwick;
llong ans[MAXN];
void buildDFS(int node, int inc, int par)
{
in[node] = ++timer;
for (const auto &[u, d] : g[node])
{
if (u == par)
{
continue;
}
buildDFS(u, d, node);
}
out[node] = timer;
SegmentTree::Node res = treeNotActive.query(in[node], out[node]);
treeNotActive.updateAdd(res.idx, inc);
};
void solveDFS(int node, int inc, int par)
{
// std::cout << "trees before: " << node << '\n';
// for (int i = 1 ; i <= n ; ++i)
// {
// std::cout << i << ": " << treeActive.query(in[i], in[i]).max << ' ' << treeActive.query(in[i], in[i]).min << ' ' << treeNotActive.query(in[i], in[i]).max << ' ' << fenwick.query(in[i], in[i]) << '\n';
// }
// std::cout << "solve dfs: " << node << ' ' << inc << ' ' << treeActive.query(1, n).sum << '\n' << std::flush;
std::vector <std::pair <int,llong>> revertActive;
std::vector <std::pair <int,llong>> revertNotActive;
std::vector <std::pair <int,int>> revertFenwick;
if (inc != 0)
{
if (fenwick.query(1, n) - fenwick.query(in[node], out[node]) != 0)
{
SegmentTree::Node outside;
if (in[node] > 1)
{
outside = treeActive.combine(outside, treeActive.query(1, in[node] - 1));
}
if (out[node] < n)
{
outside = treeActive.combine(outside, treeActive.query(out[node] + 1, n));
}
treeActive.updateAdd(outside.idx, inc);
revertActive.push_back({outside.idx, outside.max});
} else
{
SegmentTree::Node outside;
if (in[node] > 1)
{
outside = treeActive.combine(outside, treeNotActive.query(1, in[node] - 1));
}
if (out[node] < n)
{
outside = treeActive.combine(outside, treeNotActive.query(out[node] + 1, n));
}
treeNotActive.updateAdd(outside.idx, inc);
revertNotActive.push_back({outside.idx, outside.max});
}
if (fenwick.query(in[node], out[node]) != 0)
{
// std::cout << "here\n";
SegmentTree::Node inside = treeActive.query(in[node], out[node]);
int removedIdx = inside.idx;
treeActive.updateAdd(inside.idx, -inc);
revertActive.push_back({inside.idx, inside.max});
inside = treeActive.query(in[node], out[node]);
SegmentTree::Node outside;
if (in[node] > 1)
{
outside = treeActive.combine(outside, treeNotActive.query(1, in[node] - 1));
}
if (out[node] < n)
{
outside = treeActive.combine(outside, treeNotActive.query(out[node] + 1, n));
}
// if (node == 2)
// {
// std::cout << "outside max, inside min: " << outside.max << ' ' << inside.min << ' ' << in[node] << ' ' << out[node] << '\n';
// for (int i = 1 ; i <= n ; ++i)
// {
// std::cout << i << ": " << treeActive.query(in[i], in[i]).max << ' ' << treeActive.query(in[i], in[i]).min << ' ' << treeNotActive.query(in[i], in[i]).max << ' ' << fenwick.query(in[i], in[i]) << '\n';
// }
// }
if (outside.max > inside.min)
{
// std::cout << "max min: " << outside.max << ' ' << inside.min << ' ' << inside.idxMin << '\n' << std::flush;
revertFenwick.push_back({inside.idxMin, 1});
revertFenwick.push_back({outside.idx, -1});
fenwick.update(inside.idxMin, -1);
fenwick.update(outside.idx, 1);
treeActive.updateSet(inside.idxMin, -1);
treeNotActive.updateSet(inside.idxMin, inside.min);
treeActive.updateSet(outside.idx, outside.max);
treeNotActive.updateSet(outside.idx, -1);
if (removedIdx == inside.idxMin) revertActive.push_back({inside.idxMin, inside.min + inc});
else revertActive.push_back({inside.idxMin, inside.min});
revertNotActive.push_back({outside.idx, outside.max});
revertNotActive.push_back({inside.idxMin, -1});
revertActive.push_back({outside.idx, -1});
}
}
}
ans[node] = treeActive.query(1, n).sum;
// std::cout << "ans: " << node << " = " << ans[node] << '\n' << std::flush;
for (const auto &[u, d] : g[node])
{
if (u == par)
{
continue;
}
solveDFS(u, d, node);
// if (node == 1)
// {
// std::cout << "trees after: " << u << '\n';
// for (int i = 1 ; i <= n ; ++i)
// {
// std::cout << i << ": " << treeActive.query(in[i], in[i]).max << ' ' << treeNotActive.query(in[i], in[i]).max << ' ' << fenwick.query(in[i], in[i]) << '\n';
// }
// }
assert(treeActive.query(1, n).sum == ans[node]);
}
for (const auto &[pos, val] : revertActive)
{
// std::cout << "revert active: " << pos << ' ' << val << '\n';
treeActive.updateSet(pos, val);
}
for (const auto &[pos, val] : revertNotActive)
{
treeNotActive.updateSet(pos, val);
}
for (const auto &[pos, val] : revertFenwick)
{
fenwick.update(pos, val);
}
}
void solve()
{
treeActive.build(true);
treeNotActive.build(false);
buildDFS(1, 0, 0);
// std::cout << "build\n";
// for (int i = 1 ; i <= n ; ++i)
// {
// std::cout << i << ": " << treeNotActive.query(in[i], in[i]).max << '\n';
// }
int used = 0;
for (int i = 1 ; i <= k ; ++i)
{
SegmentTree::Node res = treeNotActive.query(1, n);
// std::cout << "active: " << res.idx << ' ' << res.max << '\n';
treeNotActive.updateSet(res.idx, -1);
treeActive.updateSet(res.idx, res.max);
fenwick.update(res.idx, 1);
}
solveDFS(1, 0, 0);
}
void print()
{
for (int i = 1 ; i <= n ; ++i)
{
std::cout << ans[i] << '\n';
}
}
void input()
{
std::cin >> n >> k;
for (int i = 1 ; i < n ; ++i)
{
int u, v, d;
std::cin >> u >> v >> d;
g[u].push_back({v, d});
g[v].push_back({u, d});
}
}
void fastIOI()
{
std::ios_base :: sync_with_stdio(0);
std::cout.tie(nullptr);
std::cin.tie(nullptr);
}
signed main()
{
fastIOI();
input();
solve();
print();
return 0;
}
Compilation message
Main.cpp: In member function 'void SegmentTree::build(long long int, long long int, long long int, bool)':
Main.cpp:74:21: warning: suggest parentheses around '+' inside '>>' [-Wparentheses]
74 | int mid = l + r >> 1;
| ~~^~~
Main.cpp: In member function 'void SegmentTree::updateSet(long long int, long long int, long long int, long long int, llong)':
Main.cpp:99:21: warning: suggest parentheses around '+' inside '>>' [-Wparentheses]
99 | int mid = l + r >> 1;
| ~~^~~
Main.cpp: In member function 'void SegmentTree::updateAdd(long long int, long long int, long long int, long long int, llong)':
Main.cpp:115:21: warning: suggest parentheses around '+' inside '>>' [-Wparentheses]
115 | int mid = l + r >> 1;
| ~~^~~
Main.cpp: In member function 'SegmentTree::Node SegmentTree::query(long long int, long long int, long long int, long long int, long long int)':
Main.cpp:128:21: warning: suggest parentheses around '+' inside '>>' [-Wparentheses]
128 | int mid = l + r >> 1;
| ~~^~~
Main.cpp: In function 'void solve()':
Main.cpp:362:9: warning: unused variable 'used' [-Wunused-variable]
362 | int used = 0;
| ^~~~
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
6 ms |
35416 KB |
Output is correct |
2 |
Correct |
5 ms |
35672 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
6 ms |
35416 KB |
Output is correct |
2 |
Correct |
5 ms |
35672 KB |
Output is correct |
3 |
Correct |
6 ms |
35420 KB |
Output is correct |
4 |
Correct |
6 ms |
35564 KB |
Output is correct |
5 |
Correct |
6 ms |
35504 KB |
Output is correct |
6 |
Correct |
6 ms |
35420 KB |
Output is correct |
7 |
Correct |
5 ms |
35420 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
6 ms |
35416 KB |
Output is correct |
2 |
Correct |
5 ms |
35672 KB |
Output is correct |
3 |
Correct |
6 ms |
35420 KB |
Output is correct |
4 |
Correct |
6 ms |
35564 KB |
Output is correct |
5 |
Correct |
6 ms |
35504 KB |
Output is correct |
6 |
Correct |
6 ms |
35420 KB |
Output is correct |
7 |
Correct |
5 ms |
35420 KB |
Output is correct |
8 |
Correct |
7 ms |
35416 KB |
Output is correct |
9 |
Correct |
7 ms |
35932 KB |
Output is correct |
10 |
Correct |
7 ms |
35500 KB |
Output is correct |
11 |
Correct |
6 ms |
35420 KB |
Output is correct |
12 |
Correct |
7 ms |
35528 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
6 ms |
35416 KB |
Output is correct |
2 |
Correct |
5 ms |
35672 KB |
Output is correct |
3 |
Correct |
6 ms |
35420 KB |
Output is correct |
4 |
Correct |
6 ms |
35564 KB |
Output is correct |
5 |
Correct |
6 ms |
35504 KB |
Output is correct |
6 |
Correct |
6 ms |
35420 KB |
Output is correct |
7 |
Correct |
5 ms |
35420 KB |
Output is correct |
8 |
Correct |
7 ms |
35416 KB |
Output is correct |
9 |
Correct |
7 ms |
35932 KB |
Output is correct |
10 |
Correct |
7 ms |
35500 KB |
Output is correct |
11 |
Correct |
6 ms |
35420 KB |
Output is correct |
12 |
Correct |
7 ms |
35528 KB |
Output is correct |
13 |
Correct |
9 ms |
35672 KB |
Output is correct |
14 |
Correct |
9 ms |
36444 KB |
Output is correct |
15 |
Correct |
8 ms |
35932 KB |
Output is correct |
16 |
Correct |
8 ms |
35676 KB |
Output is correct |
17 |
Correct |
9 ms |
35800 KB |
Output is correct |
18 |
Correct |
8 ms |
35676 KB |
Output is correct |
19 |
Correct |
10 ms |
35672 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
159 ms |
43436 KB |
Output is correct |
2 |
Correct |
186 ms |
59188 KB |
Output is correct |
3 |
Correct |
123 ms |
45704 KB |
Output is correct |
4 |
Correct |
137 ms |
45480 KB |
Output is correct |
5 |
Correct |
167 ms |
50276 KB |
Output is correct |
6 |
Correct |
138 ms |
45256 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
6 ms |
35416 KB |
Output is correct |
2 |
Correct |
5 ms |
35672 KB |
Output is correct |
3 |
Correct |
6 ms |
35420 KB |
Output is correct |
4 |
Correct |
6 ms |
35564 KB |
Output is correct |
5 |
Correct |
6 ms |
35504 KB |
Output is correct |
6 |
Correct |
6 ms |
35420 KB |
Output is correct |
7 |
Correct |
5 ms |
35420 KB |
Output is correct |
8 |
Correct |
7 ms |
35416 KB |
Output is correct |
9 |
Correct |
7 ms |
35932 KB |
Output is correct |
10 |
Correct |
7 ms |
35500 KB |
Output is correct |
11 |
Correct |
6 ms |
35420 KB |
Output is correct |
12 |
Correct |
7 ms |
35528 KB |
Output is correct |
13 |
Correct |
9 ms |
35672 KB |
Output is correct |
14 |
Correct |
9 ms |
36444 KB |
Output is correct |
15 |
Correct |
8 ms |
35932 KB |
Output is correct |
16 |
Correct |
8 ms |
35676 KB |
Output is correct |
17 |
Correct |
9 ms |
35800 KB |
Output is correct |
18 |
Correct |
8 ms |
35676 KB |
Output is correct |
19 |
Correct |
10 ms |
35672 KB |
Output is correct |
20 |
Correct |
159 ms |
43436 KB |
Output is correct |
21 |
Correct |
186 ms |
59188 KB |
Output is correct |
22 |
Correct |
123 ms |
45704 KB |
Output is correct |
23 |
Correct |
137 ms |
45480 KB |
Output is correct |
24 |
Correct |
167 ms |
50276 KB |
Output is correct |
25 |
Correct |
138 ms |
45256 KB |
Output is correct |
26 |
Correct |
259 ms |
45864 KB |
Output is correct |
27 |
Correct |
208 ms |
59204 KB |
Output is correct |
28 |
Correct |
250 ms |
62392 KB |
Output is correct |
29 |
Correct |
136 ms |
45692 KB |
Output is correct |
30 |
Correct |
245 ms |
45768 KB |
Output is correct |
31 |
Correct |
182 ms |
45292 KB |
Output is correct |
32 |
Correct |
292 ms |
51796 KB |
Output is correct |
33 |
Correct |
189 ms |
45908 KB |
Output is correct |
34 |
Correct |
133 ms |
46160 KB |
Output is correct |
35 |
Correct |
224 ms |
45864 KB |
Output is correct |
36 |
Correct |
260 ms |
67724 KB |
Output is correct |