Submission #1132538

#	Time	Username	Problem	Language	Result	Execution time	Memory
1132538		NK_	Paths (RMI21_paths)	C++20	100 / 100	204 ms	31220 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});
	}

	int R = 0;
	for(int i = 0; i < N; i++) if (sz(adj[i]) > 1) R = 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) {
		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 (!sz(chd[u])) return u;
		return (int) (*rbegin(chd[u]))[1];
	};	

	gen(R, R);

	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.ins(*begin(take));
			ans -= (*begin(take))[0]; 
			take.erase(begin(take)); 

			take.ins(*rbegin(have)); 
			ans += (*rbegin(have))[0];
			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;

		// ll res = 0;
		// for(auto& x : take) res += x[0];
			
		// assert(res == 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(R, R);

	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: 12 / 12
• Subtask 6: 32 / 32