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<bits/stdc++.h>
#define int long long
#define pb push_back
#define fi first
#define se second
using namespace std;
using ll = long long;
using ld = long double;
using ull = unsigned long long;
mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
const int maxN = 1e6 + 5;
const int mod = 1e9 + 7;
const ll oo = 1e18;
int n, k;
vector<pair<int, int>> adj[maxN];
int f[maxN], dp[maxN];
void ReadInput()
{
cin >> n >> k;
for(int i=1; i<n; i++)
{
int u, v, w;
cin >> u >> v >> w;
adj[u].pb({v, w});
adj[v].pb({u, w});
}
}
void predfs(int u, int par)
{
for(auto a : adj[u])
{
int v = a.fi, w = a.se;
if(v == par) continue;
predfs(v, u);
f[u] = max(f[u], f[v] + w);
}
}
void predfs1(int u, int par)
{
vector<pair<int, int>> vc;
for(auto a : adj[u])
{
int v = a.fi, w = a.se;
if(v == par) continue;
vc.pb({f[v] + w, v});
}
sort(vc.begin(), vc.end(), greater<pair<int, int>>());
for(auto a : adj[u])
{
int v = a.fi, w = a.se;
if(v == par) continue;
if(v != vc[0].se) dp[v] = max(dp[u] + w, vc[0].fi + w);
else dp[v] = max(dp[u] + w, vc[1].fi + w);
predfs1(v, u);
}
}
multiset<int, greater<int>> s, s1;
void dfs(int u, int par, int root)
{
int _max = 0, id = 0;
for(auto a : adj[u])
{
int v = a.fi, w = a.se;
if(v == par) continue;
dfs(v, u, root);
if(f[v] + w > _max)
{
_max = f[v] + w;
id = v;
}
}
// cout << u << " " << _max << "\n";
for(auto a : adj[u])
{
int v = a.fi, w = a.se;
if(v == par || v == id) continue;
s.insert(f[v] + w);
}
if(u == root) s.insert(f[u]);
}
int res[maxN], sum = 0;
void add(int val)
{
if(s.size() < k)
{
s.insert(val);
sum += val;
return;
}
auto it = s.end();
it--;
if(*it < val)
{
sum -= *it;
sum += val;
s1.insert(*it);
s.erase(it);
s.insert(val);
}
else s1.insert(val);
}
void del(int val)
{
if(s.find(val) != s.end())
{
sum -= val;
s.erase(s.find(val));
if(s1.empty()) return;
sum += *s1.begin();
s.insert(*s1.begin());
s1.erase(s1.begin());
}
else if(s1.find(val) != s1.end()) s1.erase(s1.find(val));
}
void reroot(int u, int par)
{
res[u] = sum;
for(auto a : adj[u])
{
int v = a.fi, w = a.se;
if(v == par) continue;
add(f[v]);
del(f[v] + w);
add(dp[v]);
del(dp[v] - w);
reroot(v, u);
add(f[v] + w);
del(f[v]);
add(dp[v] - w);
del(dp[v]);
}
}
void Solve()
{
predfs(n, 0);
predfs1(n, 0);
dfs(n, 0, n);
for(int v : s) sum += v;
while(s.size() > k)
{
auto it = s.end();
it--;
sum -= *it;
s1.insert(*it);
s.erase(it);
}
//for(int v : s1) cout << v << " ";return;
reroot(n, 0);
for(int i=1; i<=n; i++) cout << res[i] << '\n';
}
int32_t main()
{
// freopen("x.inp", "r", stdin);
ios_base::sync_with_stdio(false);
cin.tie(nullptr);
ReadInput();
Solve();
}
Compilation message (stderr)
Main.cpp: In function 'void add(long long int)':
Main.cpp:85:17: warning: comparison of integer expressions of different signedness: 'std::multiset<long long int, std::greater<long long int> >::size_type' {aka 'long unsigned int'} and 'long long int' [-Wsign-compare]
85 | if(s.size() < k)
| ~~~~~~~~~^~~
Main.cpp: In function 'void Solve()':
Main.cpp:140:20: warning: comparison of integer expressions of different signedness: 'std::multiset<long long int, std::greater<long long int> >::size_type' {aka 'long unsigned int'} and 'long long int' [-Wsign-compare]
140 | while(s.size() > k)
| ~~~~~~~~~^~~
# | 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... |