Submission #1132518

#	Time	Username	Problem	Language	Result	Execution time	Memory
1132518		NK_	Paths (RMI21_paths)	C++17	56 / 100	208 ms	29912 KiB

Main

// 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<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;
		}

		return best;
	};	

	gen(0, 0);

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

		while(sz(take) && sz(have) && *begin(take) < *rbegin(have)) {
			have.insert(*begin(take));
			ans -= (*begin(take))[0]; take.erase(begin(take)); 
			take.insert(*rbegin(have)); ans += (*rbegin(have))[0];
			have.erase(*prev(end(have))); 
		}
	};

	auto upd = [&](int x, ll c) {
		if (have.count({C[x], x})) have.erase({C[x], x});
		else { take.erase({C[x], x}); ans -= C[x]; }

		C[x] += c;

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

		norm();
	};

	for(int i = 0; i < N; i++) if (sz(adj[i]) == 1) {
		have.insert({C[i], i});
		chd[i].insert({0, i});
	}
	norm();

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

		for(auto& [v, w, i] : adj[u]) if (v != p) {
			int vx = (*rbegin(chd[v]))[1];
			chd[u].erase({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].insert({C[ux], ux});

			answer(v, u);

			chd[v].erase({C[ux], ux});
			upd(ux, -w);

			upd(vx, +w);
			chd[u].insert({C[vx], vx});
		}
	};

	answer(0, 0);

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

	exit(0-0);
}	

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