#include <bits/stdc++.h>
#define pb push_back
#define int int64_t
#define pii pair<int, int>
#define X first
#define Y second
#define all(x) (x).begin(),(x).end()
using namespace std;
const int inf = 1e18;
struct edge{
int v, w;
};
struct bag{
int cur, k;
multiset<int> lower, upper;
bag(){
//cout << "Hello";
cur = 0;
k = 0;
lower.clear();
upper.clear();
}
void balance(){
if(lower.empty())return;
while(lower.size() && upper.size() < k){
cur+=*lower.rbegin();
upper.insert(*lower.rbegin());
auto it = lower.end(); it--;
lower.erase(it);
}
if(lower.empty())return;
while(*lower.rbegin() > *upper.begin()){
//cout << "start balance \n";
//p();
auto tu = upper.begin();
auto tl = lower.end(); tl--;
int vtu = *tu;
int vtl = *tl;
cur-= vtu;
cur+= vtl;
upper.erase(tu);
lower.erase(tl);
upper.insert(vtl);
lower.insert(vtu);
//cout << "end balance \n";
//p();
}
}
void insert(int a){
//cout << "ins: " << a << '\n';
lower.insert(a);
balance();
}
void remove(int a){
if(lower.find(a) != lower.end()){
lower.erase(lower.lower_bound(a));
}else{
upper.erase(upper.lower_bound(a));
cur-=a;
balance();
}
}
void p(){
/*
cout << cur << '\n';
cout << "lower: ";
for(auto e : lower)cout << e << ' ';
cout << '\n';
cout << "upper: ";
for(auto e : upper)cout << e << ' ';
cout << '\n';
*/
}
};
vector<vector<edge>> g;
vector<int> up;
vector<int> down;
vector<int> ans;
void merge(multiset<int>& A, multiset<int>& B){
if(A.size() < B.size())A.swap(B);
for(auto e : B)A.insert(e);
}
multiset<int> dfs1(int u, int p, int lvl){
up[u] = lvl;
multiset<int> S;
S.insert(lvl);
for(auto [v, w] : g[u]){
if(v == p)continue;
multiset<int> tmp = dfs1(v, u, lvl+w);
if(*tmp.rbegin() > *S.rbegin())S.swap(tmp);
auto it = tmp.end(); it--;
tmp.erase(it);
tmp.insert(*it-lvl);
merge(S, tmp);
}
down[u] = *S.rbegin() - lvl;
return S;
}
void dfs2(int u, int p, bag& B){
//cout << "start " << u << '\n';
B.p();
ans[u] = B.cur;
vector<int> mx;
mx.pb(0);
for(auto [v, w] : g[u])mx.pb(down[v] + w);
sort(all(mx));
reverse(all(mx));
for(auto [v, w] : g[u]){
if(v == p)continue;
int prv_down = down[u];
if(down[v] + w == mx[0]){
down[u] = mx[1];
B.remove(mx[1]); //cout << "rem: " << mx[1] << '\n';
B.insert(mx[1] + w); //cout << "ins: " << (mx[1] + w) << '\n';
B.remove(mx[0]); //cout << "rem: " << mx[0] << '\n';
B.insert(mx[0] - w); //cout << "ins: " << (mx[0] - w) << '\n';
dfs2(v, u, B);
B.p();
down[u] = prv_down;
B.remove(mx[1] + w);
B.insert(mx[1]);
B.remove(mx[0] - w);
B.insert(mx[0]);
}else{
down[u] = mx[0];
B.remove(mx[0]); //cout << "rem: " << mx[0] << '\n';
B.insert(mx[0] + w); //cout << "ins: " << (mx[0] + w) << '\n';
B.remove(down[v] + w); //cout << "rem: " << (down[v] + w) << '\n';
B.insert(down[v]); //cout << "ins: " << (down[v]) << '\n';
dfs2(v, u, B);
down[u] = prv_down;
B.remove(mx[0] + w);
B.insert(mx[0]);
B.remove(down[v]);
B.insert(down[v] + w);
}
}
//cout <<"fin " << u << '\n';
}
void solve(){
int n, k; cin >> n >> k;
g.assign(n, vector<edge>());
up.assign(n, 0);
down.assign(n, 0);
ans.assign(n, 0);
for(int i = 0; i< n-1; i++){
int c, d, w; cin >> c >> d >> w; c--; d--;
g[c].pb({d, w});
g[d].pb({c, w});
}
multiset<int> S = dfs1(0, 0, 0);
//for(auto e : S)cout << e << ' '; cout << '\n';
int cur = 0;
bag B; B.k = k; B.cur = 0;
for(auto e : S) B.insert(e);
B.p();
//cout << B.cur << '\n';
dfs2(0, 0, B);
for(auto e : ans)cout << e << '\n';
}
signed main(){
//ios_base::sync_with_stdio(false);
//cin.tie(NULL);
int t = 1;
//cin >> t;
while(t--)solve();
}
Compilation message
Main.cpp: In member function 'void bag::balance()':
Main.cpp:30:38: warning: comparison of integer expressions of different signedness: 'std::multiset<long int>::size_type' {aka 'long unsigned int'} and 'int64_t' {aka 'long int'} [-Wsign-compare]
30 | while(lower.size() && upper.size() < k){
| ~~~~~~~~~~~~~^~~
Main.cpp: In function 'void solve()':
Main.cpp:172:6: warning: unused variable 'cur' [-Wunused-variable]
172 | int cur = 0;
| ^~~
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
0 ms |
204 KB |
Output is correct |
2 |
Correct |
0 ms |
204 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
0 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 |
2 ms |
332 KB |
Output is correct |
5 |
Correct |
1 ms |
332 KB |
Output is correct |
6 |
Correct |
1 ms |
332 KB |
Output is correct |
7 |
Correct |
1 ms |
204 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
0 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 |
2 ms |
332 KB |
Output is correct |
5 |
Correct |
1 ms |
332 KB |
Output is correct |
6 |
Correct |
1 ms |
332 KB |
Output is correct |
7 |
Correct |
1 ms |
204 KB |
Output is correct |
8 |
Correct |
3 ms |
460 KB |
Output is correct |
9 |
Correct |
3 ms |
588 KB |
Output is correct |
10 |
Correct |
3 ms |
460 KB |
Output is correct |
11 |
Correct |
4 ms |
480 KB |
Output is correct |
12 |
Correct |
3 ms |
460 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
0 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 |
2 ms |
332 KB |
Output is correct |
5 |
Correct |
1 ms |
332 KB |
Output is correct |
6 |
Correct |
1 ms |
332 KB |
Output is correct |
7 |
Correct |
1 ms |
204 KB |
Output is correct |
8 |
Correct |
3 ms |
460 KB |
Output is correct |
9 |
Correct |
3 ms |
588 KB |
Output is correct |
10 |
Correct |
3 ms |
460 KB |
Output is correct |
11 |
Correct |
4 ms |
480 KB |
Output is correct |
12 |
Correct |
3 ms |
460 KB |
Output is correct |
13 |
Correct |
11 ms |
588 KB |
Output is correct |
14 |
Correct |
6 ms |
844 KB |
Output is correct |
15 |
Correct |
8 ms |
716 KB |
Output is correct |
16 |
Correct |
7 ms |
588 KB |
Output is correct |
17 |
Correct |
7 ms |
716 KB |
Output is correct |
18 |
Correct |
4 ms |
588 KB |
Output is correct |
19 |
Correct |
6 ms |
588 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
410 ms |
20552 KB |
Output is correct |
2 |
Correct |
404 ms |
25820 KB |
Output is correct |
3 |
Correct |
345 ms |
20516 KB |
Output is correct |
4 |
Correct |
397 ms |
20580 KB |
Output is correct |
5 |
Correct |
407 ms |
22476 KB |
Output is correct |
6 |
Correct |
412 ms |
20560 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
0 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 |
2 ms |
332 KB |
Output is correct |
5 |
Correct |
1 ms |
332 KB |
Output is correct |
6 |
Correct |
1 ms |
332 KB |
Output is correct |
7 |
Correct |
1 ms |
204 KB |
Output is correct |
8 |
Correct |
3 ms |
460 KB |
Output is correct |
9 |
Correct |
3 ms |
588 KB |
Output is correct |
10 |
Correct |
3 ms |
460 KB |
Output is correct |
11 |
Correct |
4 ms |
480 KB |
Output is correct |
12 |
Correct |
3 ms |
460 KB |
Output is correct |
13 |
Correct |
11 ms |
588 KB |
Output is correct |
14 |
Correct |
6 ms |
844 KB |
Output is correct |
15 |
Correct |
8 ms |
716 KB |
Output is correct |
16 |
Correct |
7 ms |
588 KB |
Output is correct |
17 |
Correct |
7 ms |
716 KB |
Output is correct |
18 |
Correct |
4 ms |
588 KB |
Output is correct |
19 |
Correct |
6 ms |
588 KB |
Output is correct |
20 |
Correct |
410 ms |
20552 KB |
Output is correct |
21 |
Correct |
404 ms |
25820 KB |
Output is correct |
22 |
Correct |
345 ms |
20516 KB |
Output is correct |
23 |
Correct |
397 ms |
20580 KB |
Output is correct |
24 |
Correct |
407 ms |
22476 KB |
Output is correct |
25 |
Correct |
412 ms |
20560 KB |
Output is correct |
26 |
Correct |
531 ms |
20936 KB |
Output is correct |
27 |
Correct |
403 ms |
25636 KB |
Output is correct |
28 |
Correct |
453 ms |
26852 KB |
Output is correct |
29 |
Correct |
368 ms |
20720 KB |
Output is correct |
30 |
Correct |
428 ms |
20920 KB |
Output is correct |
31 |
Correct |
433 ms |
20712 KB |
Output is correct |
32 |
Correct |
461 ms |
23156 KB |
Output is correct |
33 |
Correct |
473 ms |
20932 KB |
Output is correct |
34 |
Correct |
374 ms |
20796 KB |
Output is correct |
35 |
Correct |
450 ms |
20988 KB |
Output is correct |
36 |
Correct |
447 ms |
28384 KB |
Output is correct |