Submission #1219940

#	Time	Username	Problem	Language	Result	Execution time	Memory
1219940		nguyn	Paths (RMI21_paths)	C++20	100 / 100	136 ms	17932 KiB

Main

#include <bits/stdc++.h>
using namespace std;

#define int long long
#define ll long long 
#define F first
#define S second
#define pb push_back 
#define pii pair<int,int>

const int N = 1e5 + 5;

int k, n; 
vector<pii> g[N]; 
multiset<int> st;
multiset<int> k_max;
int res = 0; 
int mx[N];
int ans[N];
int a[N];

void pre_dfs(int u, int p) {
	for (auto it : g[u]) {
		int v = it.F;
		int w = it.S; 
		if (v == p) continue;
		a[v] = w;
		if (g[v].size() == 1) {
			st.insert(w); 
			mx[u] = max(mx[u], w); 
			continue; 
		} 
		pre_dfs(v, u);
		if (st.find(mx[v]) != st.end()) st.erase(st.find(mx[v])); 
		st.insert(mx[v] + w); 
		mx[u] = max(mx[u], mx[v] + w); 
	}	
} 

void change(int cur, int nxt) {
	// cout << cur << ' ' << nxt << '\n';
	if (st.find(cur) != st.end()) {
		st.erase(st.find(cur)); 
		st.insert(nxt); 
	}
	else if (k_max.find(cur) != k_max.end()) {
		k_max.erase(k_max.find(cur)); 
		k_max.insert(nxt); 
		res += nxt - cur; 
	}
	else {
		st.insert(nxt);
	}
}

void dfs(int u, int p) {
	while(!st.empty()) {
		auto it = prev(st.end()); 
		if (!st.empty() && (*it) > (*k_max.begin())) {
			int cur = (*it);
			int nxt = (*k_max.begin());
			k_max.erase(k_max.find(nxt)); 
			st.erase(st.find(cur)); 
			k_max.insert(cur);
			st.insert(nxt);
			res += cur - nxt; 
		}	
		else break;
	}
	// cout << u << ":\n";
	// for (auto it : st) {
	// 	cout << it << ' ';
	// }
	// for (auto it : k_max) {
	// 	cout << it << ' ';
	// }
	// cout << endl;
	ans[u] = res;
	int fir = 0;
	int sec = 0;  
	for (auto it : g[u]) {
		int v = it.F;
		int w = it.S; 
		if (mx[v] + w > fir) {
			sec = fir; 
			fir = mx[v] + w;
		}
		else if (mx[v] + w > sec) {
			sec = mx[v] + w; 
		}
		// cout << mx[v] + w << ' '; 
	} 
	// cout << '\n'; 
	for (auto it : g[u]) {
		int v = it.F;
		int w = it.S; 
		if (v == p) continue;
		int mxu = mx[u];
		int mxv = mx[v]; 
		if (mx[v] + w == fir) {
			mx[u] = sec;
		}
		else {
			mx[u] = fir; 
		}
		// if (u == 2 && v == 3) {
		// 	cout << mx[u] << ' ' << mx[v] << ' ' << fir << '\n'; 
		// }
		int cur1 = mx[u];
		int nxt1 = mx[u] + w;
		int cur2 = mx[v] + w;
		int nxt2 = mx[v]; 
		change(cur1, nxt1); 
		change(cur2, nxt2);
		mx[v] = max(mx[v], mx[u] + w);
		dfs(v, u); 
		change(nxt1, cur1);
		change(nxt2, cur2); 
		mx[u] = mxu;
		mx[v] = mxv;  
	} 
}

signed main(){
    ios_base::sync_with_stdio(false) ; 
    cin.tie(0) ; cout.tie(0) ; 
    if (fopen("INP.INP" ,"r")) {
        freopen("INP.INP" ,"r" , stdin) ;
        freopen("OUT.OUT" , "w" , stdout) ;
    }
    cin >> n >> k;
    for (int i = 1; i < n; i++) {
    	int u, v, w;
    	cin >> u >> v >> w;
    	g[u].pb({v, w});
    	g[v].pb({u, w}); 
    }	
    pre_dfs(1, 0);
    while(k_max.size() < k) {
    	auto it = --st.end();
    	res += *it;
    	k_max.insert(*it); 
    	st.erase(it);
    }
    dfs(1, 0);
    // cout << res << '\n';
    for (int i = 1; i <= n; i++) {
    	cout << ans[i] << '\n';
    }
}

Compilation message (stderr)

Main.cpp: In function 'int main()':
Main.cpp:128:16: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
  128 |         freopen("INP.INP" ,"r" , stdin) ;
      |         ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~
Main.cpp:129:16: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
  129 |         freopen("OUT.OUT" , "w" , stdout) ;
      |         ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~

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