Submission #1132553

#TimeUsernameProblemLanguageResultExecution timeMemory
1132553NK_Paths (RMI21_paths)C++17
68 / 100
253 ms32724 KiB
// Success consists of going from failure to failure without loss of enthusiasm
#include <bits/stdc++.h>

using namespace std;

#define nl '\n'
#define pb push_back
#define sz(x) int(x.size())
#define ins insert

using ll = long long;
template<class T> using V = vector<T>;
template<class T, size_t SZ> using AR = array<T, SZ>;
template<class T> using mset = multiset<T>;
using vi = V<int>;
using vl = V<ll>;

int main() {
	cin.tie(0)->sync_with_stdio(0);
	
	int N, K; cin >> N >> K;

	V<V<AR<int, 3>>> adj(N); for(int i = 0; i < N - 1; i++) {
		int u, v, w; cin >> u >> v >> w; --u, --v;
		adj[u].pb({v, w, i});
		adj[v].pb({u, w, i});
	}

	set<int> leaf;
	for(int i = 0; i < N; i++) if (sz(adj[i]) == 1) {
		leaf.ins(i);
	}

	set<AR<ll, 2>> have, take; vl C(N);
	V<set<AR<ll, 2>>> chd(N);
	function<int(int, int)> gen = [&](int u, int p) {
		int best = u;
		for(auto& [v, w, i] : adj[u]) if (v != p) {
			int x = gen(v, u);
			C[x] += w; chd[u].insert({C[x], x});
			if (C[best] <= C[x]) best = x;
		}
		  
		if (!sz(chd[u])) return u;
		assert(C[best] == (*rbegin(chd[u]))[0]);
		assert(leaf.count(best));
		return best;
	};	

	gen(0, 0);

	ll ans = 0;
	auto norm = [&]() {
		while(sz(take) < K && sz(have)) {
			ans += (*rbegin(have))[0];
			take.ins(*rbegin(have)); 
			have.erase(prev(end(have)));
		}

		while(sz(take) && sz(have) && *begin(take) < *rbegin(have)) {
			ans -= (*begin(take))[0]; 
			ans += (*rbegin(have))[0];

			have.ins(*begin(take));
			take.erase(begin(take)); 

			take.ins(*rbegin(have)); 
			have.erase(prev(end(have))); 
		}
	};

	auto upd = [&](int x, ll c) {
		auto it = have.find(AR<ll, 2>{C[x], x});
		if (it != have.end()) have.erase(it);
		else { 
			// assert(take.count(AR<ll, 2>{C[x], x}));
			take.erase(AR<ll, 2>{C[x], x}); ans -= C[x]; 
		}

		C[x] += c;

		have.ins({C[x], x});
	};

	for(int i = 0; i < N; i++) if (sz(adj[i]) == 1) {
		have.ins(AR<ll, 2>{C[i], i});
		chd[i].ins(AR<ll, 2>{0, i});
	}

	// cout << sz(take) << endl;

	vl ANS(N);
	function<void(int, int)> answer = [&](int u, int p) {
		norm();
		// cout << ans << " " << (*begin(take))[0] << endl;
		// ANS[u] = (*begin(take))[0];
		ANS[u] = ans;

		for(auto& [v, w, i] : adj[u]) if (v != p) {
			int vx = (*rbegin(chd[v]))[1];
			chd[u].erase(AR<ll, 2>{C[vx], vx});
			upd(vx, -w); // update by -w (remove edge from vx)

			int ux = (*rbegin(chd[u]))[1];
			upd(ux, +w); // give edge to best -> ux
			chd[v].ins(AR<ll, 2>{C[ux], ux});

			answer(v, u);

			chd[v].erase(AR<ll, 2>{C[ux], ux});
			upd(ux, -w);

			upd(vx, +w);
			chd[u].ins(AR<ll, 2>{C[vx], vx});
		}
	};

	answer(0, 0);

	for(auto& x : ANS) cout << x << nl;

	exit(0-0);
}	
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...