Submission #514770

#	Time	Username	Problem	Language	Result	Execution time	Memory
514770		kshitij_sodani	Paths (RMI21_paths)	C++14	100 / 100	427 ms	29416 KiB

Main

//#pragma GCC optimize("Ofast,unroll-loops")
#include <bits/stdc++.h>
using namespace std;
typedef long long llo;
#define mp make_pair
#define pb push_back
#define a first 
#define b second
#define endl '\n'
llo n,k;
vector<pair<llo,llo>> adj[100001];
llo ans[100001];
pair<llo,llo> val[100001];
pair<llo,llo> val2[100001];

#include <ext/pb_ds/tree_policy.hpp>
#include <ext/pb_ds/assoc_container.hpp>
using namespace __gnu_pbds;
#define ord tree<pair<llo,llo>,null_type,greater<pair<llo,llo>>,rb_tree_tag,tree_order_statistics_node_update>
ord cur;
//vector<llo> cur2;
vector<pair<llo,llo>> pre[100001];
void dfs(llo no,llo par=-1){
	vector<pair<pair<llo,llo>,llo>> ss;

	for(auto j:adj[no]){
		if(j.a!=par){
			
			dfs(j.a,no);
			ss.pb({{val[j.a].a+j.b,val[j.a].b},j.a});
		}
	}
	if(ss.size()==0){
		val[no]={0,no};

		return;
	}
	sort(ss.begin(),ss.end());
	for(llo j=0;j<ss.size();j++){
		pre[no].pb({ss[j].a.a,ss[j].b});
		if(j+1==ss.size()){
			val[no]=ss[j].a;
		}
		else{

			//cur2.pb(ss[j].a);
			cur.insert(ss[j].a);
		}
	}
}
llo su=0;
vector<pair<pair<llo,llo>,llo>> rr;

void remove(pair<llo,llo> x){
	if(cur.find(x)==cur.end()){
		return;
	}

	//rr.pb({x,0});
	llo xx=cur.order_of_key(x);
	if(xx<k){
		su-=x.a;
		cur.erase(x);
		if(cur.size()>=k){
			su+=((*cur.find_by_order(k-1)).a);
		}
	}
	else{
		cur.erase(x);
	}
}

void add(pair<llo,llo> x){
	llo xx=cur.order_of_key(x);
	//rr.pb({x,1});
	if(cur.find(x)!=cur.end()){
		return;
	}
	if(xx<k){
		su+=x.a;
		cur.insert(x);
		if(cur.size()>=k+1){
			su-=((*cur.find_by_order(k)).a);
		}

	}
	else{
		cur.insert(x);
	}
	

}
/*void undo(){
	if(rr.back().b==1){
		remove(rr.back().a);
	}
	else{
		add(rr.back().a);
	}
	rr.pop_back();

}*/
void dfs2(llo no,llo par=-1,pair<llo,llo> ma={0,-1}){
	add(ma);
	ans[no]=su;
	/*for(llo j=0;j<k;j++){
		if(j+1>=cur.size()){
			continue;
		}
		ans[no]+=(*cur.find_by_order(j)).a;
	}*/
	/*cout<<no<<"::"<<su<<endl;
	for(auto j:cur){
		cout<<j.a<<","<<j.b<<endl;
	}
	cout<<endl;
	cout<<ma.a<<"<<"<<ma.b<<endl;*/
	remove(ma);

	vector<pair<pair<llo,llo>,llo>> ss;
	for(auto j:adj[no]){
		if(j.a!=par){
			ss.pb({{val[j.a].a+j.b,val[j.a].b},j.a});
		}
	}
	sort(ss.begin(),ss.end());
	reverse(ss.begin(),ss.end());
	for(auto j:adj[no]){
		if(j.a!=par){
			/*if(j.a==3){
				continue;
			}*/
			remove({val[j.a].a+j.b,val[j.a].b});
			add({val[j.a].a,val[j.a].b});
			pair<llo,llo> xx=val[j.a];
			/*for(auto jj:pre[no]){

			}*/
			pair<llo,llo> ma2=ma;
			llo cot=2;
			if(val2[no].b!=val[j.a].b){
				
				ma2.a+=j.b;

				ma2=max(ma2,{val2[no].a+j.b,val2[no].b});
				val[j.a]=max(val[j.a],{val[no].a+j.b,val[no].b});
				if(ma2.b==val2[no].b){
					//cout<<11111<<":"<<ma.a<<":"<<ma.b<<endl;
					remove({val2[no].a,val2[no].b});
					add({val2[no].a+j.b,val2[no].b});
					add(ma);
				}
				
				
			}
			else{
				ma2.a+=j.b;
			//	cout<<1111<<endl;
				if(ss.size()>1){
					ma2=max(ma2,{ss[1].a.a+j.b,ss[1].a.b});
					if(ma2.b==ss[1].a.b){
						remove(ss[1].a);
						cot++;
						add(ma);
						cot++;
					}
					else{

					}
				}
			}
			dfs2(j.a,no,ma2);
			
			val[j.a]=xx;
			if(val2[no].b!=val[j.a].b){
				if(ma2.b==val2[no].b){
					remove(ma);
					remove({val2[no].a+j.b,val2[no].b});
					add({val2[no].a,val2[no].b});
				}
			}
			else{
				if(ss.size()>1){
					if(ma2.b==ss[1].a.b){

						remove(ma);
						add(ss[1].a);
					}
					else{

					}
				}
			}
			remove({val[j.a].a,val[j.a].b});
			add({val[j.a].a+j.b,val[j.a].b});
			





		}
	}
	//undo();
}
int main(){
	ios_base::sync_with_stdio(false);
	cin.tie(NULL);
	cin>>n>>k;
	llo ma=0;
	for(llo i=0;i<n-1;i++){
		llo aa,bb,cc;
		cin>>aa>>bb>>cc;
		aa--;
		bb--;
		ma=cc;
		adj[aa].pb({bb,cc});
		adj[bb].pb({aa,cc});
	}
	
	if(n==1){
		cout<<0<<endl;
		return 0;
	}
	if(n==2){
		cout<<ma<<endl;
		cout<<ma<<endl;
		return 0;
	}
	/*for(llo i=0;i<n;i++){
		cur2.clear();
		dfs(i);
		llo ans2=0;
		//cur.insert(val[i]);
		cur2.pb(val[i].a);
		sort(cur2.begin(),cur2.end());
		reverse(cur2.begin(),cur2.end());
		for(llo j=0;j<k;j++){
			if(j>=cur2.size()){
				continue;
			}
			ans2+=cur2[j];
		}
		cout<<ans2<<endl;
	}*/
	llo ind=0;
	for(llo i=0;i<n;i++){
		if(adj[i].size()>1){
			ind=i;
		}
	}
	dfs(ind);
	for(llo i=0;i<n;i++){
		val2[i]=val[i];
	}
	k=min(k,(llo)(cur.size()));
	
	cur.insert(val[ind]);
	for(llo j=0;j<k;j++){
		su+=(*cur.find_by_order(j)).a;
	}
/*	for(auto j:cur){
		cout<<j.a<<":"<<j.b<<endl;
	}
	cout<<ind<<":"<<su<<endl;*/
	dfs2(ind);
	for(llo i=0;i<n;i++){
		cout<<ans[i]<<endl;
	}



 
	return 0;
}

Compilation message (stderr)

Main.cpp: In function 'void dfs(llo, llo)':
Main.cpp:39:15: warning: comparison of integer expressions of different signedness: 'llo' {aka 'long long int'} and 'std::vector<std::pair<std::pair<long long int, long long int>, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   39 |  for(llo j=0;j<ss.size();j++){
      |              ~^~~~~~~~~~
Main.cpp:41:9: warning: comparison of integer expressions of different signedness: 'llo' {aka 'long long int'} and 'std::vector<std::pair<std::pair<long long int, long long int>, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   41 |   if(j+1==ss.size()){
      |      ~~~^~~~~~~~~~~
Main.cpp: In function 'void remove(std::pair<long long int, long long int>)':
Main.cpp:64:16: warning: comparison of integer expressions of different signedness: '__gnu_pbds::detail::bin_search_tree_set<std::pair<long long int, long long int>, __gnu_pbds::null_type, std::greater<std::pair<long long int, long long int> >, __gnu_pbds::detail::tree_traits<std::pair<long long int, long long int>, __gnu_pbds::null_type, std::greater<std::pair<long long int, long long int> >, __gnu_pbds::tree_order_statistics_node_update, __gnu_pbds::rb_tree_tag, std::allocator<char> >, std::allocator<char> >::size_type' {aka 'long unsigned int'} and 'llo' {aka 'long long int'} [-Wsign-compare]
   64 |   if(cur.size()>=k){
      |      ~~~~~~~~~~^~~
Main.cpp: In function 'void add(std::pair<long long int, long long int>)':
Main.cpp:82:16: warning: comparison of integer expressions of different signedness: '__gnu_pbds::detail::bin_search_tree_set<std::pair<long long int, long long int>, __gnu_pbds::null_type, std::greater<std::pair<long long int, long long int> >, __gnu_pbds::detail::tree_traits<std::pair<long long int, long long int>, __gnu_pbds::null_type, std::greater<std::pair<long long int, long long int> >, __gnu_pbds::tree_order_statistics_node_update, __gnu_pbds::rb_tree_tag, std::allocator<char> >, std::allocator<char> >::size_type' {aka 'long unsigned int'} and 'llo' {aka 'long long int'} [-Wsign-compare]
   82 |   if(cur.size()>=k+1){
      |      ~~~~~~~~~~^~~~~

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