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>
using namespace std;
#define fi first
#define se second
#define mp make_pair
#define pb emplace_back
#define ALL(x) x.begin(),x.end()
#define SZ(x) (int)x.size()
typedef long long ll;
typedef pair<int,int> ii;
typedef pair<ii,ii> i4;
const int MOD=1000000007;
const int INF=1012345678;
const ll LLINF=1012345678012345678LL;
const double PI=3.1415926536;
const double EPS=1e-14;
const int MAXN=100005;
// problem paths: choose k vertices and take sum of all the vertices
// that you visit by travelling from root to those vertices
// across all possible roots
int n,k;
vector<ii> adj[MAXN];
multiset<ll> take,notake; // (path-to-leaf)s. take stores the top <=k
ll ans[MAXN],cur=0;
ll path[MAXN]; // the best path for this subtree (subtree based on depth from new root)
void add(ll x){
take.insert(x);
cur+=x;
if(SZ(take)>k){
ll y=(*take.begin());
take.erase(take.begin()); // multiset
cur-=y;
notake.insert(y);
}
}
void rem(ll x){
if(take.find(x)!=take.end()){
take.erase(take.find(x));
cur-=x;
while(!notake.empty()&&SZ(take)<k){
ll y=(*(--notake.end()));
take.insert(y);
cur+=y;
notake.erase(--notake.end());
}
}else{
if(notake.find(x)!=notake.end())notake.erase(notake.find(x));
}
}
void dfs(int x,int p){
path[x]=0;
if(SZ(adj[x])==1&&p!=0){ // is a leaf
add(0);
return;
}
for(ii i:adj[x]){
if(i.fi==p)continue;
dfs(i.fi,x);
rem(path[i.fi]);
add(path[i.fi]+i.se);
path[x]=max(path[x],path[i.fi]+i.se);
// only the best path that starts from x will progress to p
// similarly all the paths that started from i.fi stay there,
// except for the best one that progresses to x
// (progress -> can consider in take/notake)
}
}
void reroot(int x,int p){
ans[x]=cur; // take/notake currently considers x as the root
pair<ll,int> best,sec;
for(ii i:adj[x]){ // finding the longest and second longest paths leading to the current root
if(path[i.fi]+i.se>=best.fi){
sec=best;
best=mp(path[i.fi]+i.se,i.fi);
}else if(path[i.fi]+i.se>sec.fi){
sec=mp(path[i.fi]+i.se,i.fi);
}
}
ll old=path[x];
for(ii i:adj[x]){ // try all children as possible roots
if(i.fi==p)continue;
if(best.se==i.fi){ // this child was the best so we will extend the second best
path[x]=sec.fi;
}else{ // extend the best
path[x]=best.fi;
}
rem(path[x]); add(path[x]+i.se);
rem(path[i.fi]+i.se); add(path[i.fi]);
reroot(i.fi,x);
rem(path[i.fi]); add(path[i.fi]+i.se);
rem(path[x]+i.se); add(path[x]);
path[x]=old;
}
}
int main(){
scanf("%d%d",&n,&k);
if(n==1){printf("0\n");return 0;}
int a,b,c;
for(int i=1;i<n;i++){
scanf("%d%d%d",&a,&b,&c);
adj[a].pb(b,c);
adj[b].pb(a,c);
}
dfs(1,0);
reroot(1,0);
for(int i=1;i<=n;i++)printf("%lld\n",ans[i]);
}
Compilation message (stderr)
Main.cpp: In function 'int main()':
Main.cpp:103:7: warning: ignoring return value of 'int scanf(const char*, ...)' declared with attribute 'warn_unused_result' [-Wunused-result]
103 | scanf("%d%d",&n,&k);
| ~~~~~^~~~~~~~~~~~~~
Main.cpp:107:8: warning: ignoring return value of 'int scanf(const char*, ...)' declared with attribute 'warn_unused_result' [-Wunused-result]
107 | scanf("%d%d%d",&a,&b,&c);
| ~~~~~^~~~~~~~~~~~~~~~~~~
# | 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... |