Submission #538948

#	Time	Username	Problem	Language	Result	Execution time	Memory
538948		jamielim	Paths (RMI21_paths)	C++14	100 / 100	357 ms	16052 KiB

Main

#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);
      |   ~~~~~^~~~~~~~~~~~~~~~~~~

• Subtask 1: 8 / 8
• Subtask 2: 11 / 11
• Subtask 3: 17 / 17
• Subtask 4: 20 / 20
• Subtask 5: 12 / 12
• Subtask 6: 32 / 32