This submission is migrated from previous version of oj.uz, which used different machine for grading. This submission may have different result if resubmitted.
#include <algorithm>
#include <iostream>
#include <numeric>
#include <cassert>
#include <vector>
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::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)
{
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 (outside.max > inside.min)
{
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;
for (const auto &[u, d] : g[node])
{
if (u == par)
{
continue;
}
solveDFS(u, d, node);
}
for (const auto &[pos, val] : revertActive)
{
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);
int used = 0;
for (int i = 1 ; i <= k ; ++i)
{
SegmentTree::Node res = treeNotActive.query(1, 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);
}
int main()
{
fastIOI();
input();
solve();
print();
return 0;
}
Compilation message (stderr)
Main.cpp: In member function 'void SegmentTree::build(int, int, int, bool)':
Main.cpp:73:21: warning: suggest parentheses around '+' inside '>>' [-Wparentheses]
73 | int mid = l + r >> 1;
| ~~^~~
Main.cpp: In member function 'void SegmentTree::updateSet(int, int, int, int, llong)':
Main.cpp:98:21: warning: suggest parentheses around '+' inside '>>' [-Wparentheses]
98 | int mid = l + r >> 1;
| ~~^~~
Main.cpp: In member function 'void SegmentTree::updateAdd(int, int, int, int, llong)':
Main.cpp:114:21: warning: suggest parentheses around '+' inside '>>' [-Wparentheses]
114 | int mid = l + r >> 1;
| ~~^~~
Main.cpp: In member function 'SegmentTree::Node SegmentTree::query(int, int, int, int, int)':
Main.cpp:127:21: warning: suggest parentheses around '+' inside '>>' [-Wparentheses]
127 | int mid = l + r >> 1;
| ~~^~~
Main.cpp: In function 'void solve()':
Main.cpp:325:9: warning: unused variable 'used' [-Wunused-variable]
325 | int used = 0;
| ^~~~
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |