Submission #499826

#	Time	Username	Problem	Language	Result	Execution time	Memory
499826		Meloric	Paths (RMI21_paths)	C++17	100 / 100	531 ms	28384 KiB

Main

#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 (stderr)

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

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