제출 #1132496

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
1132496		NK_	Paths (RMI21_paths)	C++17	56 / 100	1100 ms	167284 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>;

const ll MX = 1LL << 47;
unordered_map<ll, ll> AMT, SUM; ll all = 0;
void add(ll p, ll x) { 
	// cout << "ADD: " << p << " " << x << endl;
	ll xi = p * x; all += xi;
	for(++p; p <= MX; p += p & -p) { AMT[p - 1] += x; SUM[p - 1] += xi; }	
}

ll get(ll k) { // sum of last k
	if (k == 0) return 0;
	ll p = 0, ans = 0, have = -MX;
	for(ll x = MX >> 1; x; x >>= 1) {
		const ll &amt = AMT[p + x - 1];
		if (amt <= k && p + x <= MX) {
			k -= amt;
			ans += SUM[p + x - 1];
			p += x;
		}

		if (amt) have = p + x - 1;
	}

	ans += k * have; // extra room
	return ans;
}

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

	V<V<AR<int, 2>>> 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});
		adj[v].pb({u, w});
	}

	V<mset<ll>> chd(N); vl todo;
	function<ll(int, int)> gen = [&](int u, int p) {
		for(auto& [v, w] : adj[u]) if (v != p) {
			ll x = gen(v, u) + w;
			chd[u].ins(x);
		}

		int cur = 0;
		for(auto& x : chd[u]) {
			++cur;
			if (cur == sz(chd[u])) break;
			todo.pb(x);
		}

		if (!sz(chd[u])) return 0LL;
		return *rbegin(chd[u]);
	};	

	todo.pb(gen(0, 0));

	if (sz(adj[0]) == 1) add(0, +1);

	// MX = 1LL << __lg(*max_element(begin(todo), end(todo))); MX *= 2;
	for(auto& x : todo) add(x, +1);

	int L = 0; for(int i = 0; i < N; i++) if (sz(adj[i]) == 1) {
		++L;
		chd[i].insert(0);
	}
	// cout << "LEAF: " << L << endl;

	vl ans(N);
	function<void(int, int)> answer = [&](int u, int p) {
		ans[u] = all - get(max(0, L - K));
		// cout << L - K << endl;
		// cout << u << " ==> " << ans[u] << endl;

		for(auto& [v, w] : adj[u]) if (v != p) {
			// cout << sz(chd[v]) << endl;
			ll mxv = *rbegin(chd[v]) + w;
			chd[u].erase(mxv); add(mxv, -1); // rem v as chd of u
			add(*rbegin(chd[v]), +1);

			// cout << sz(chd[u]) << endl;
			add(*rbegin(chd[u]), -1);
			ll mxu = *rbegin(chd[u]) + w;
			chd[v].insert(mxu); add(mxu, +1); // add u as chd of v

			answer(v, u);

			chd[v].erase(mxu); add(mxu, -1); // rem u as chd of v
			add(*rbegin(chd[u]), +1);

			add(*rbegin(chd[v]), -1);
			chd[u].insert(mxv); add(mxv, +1); // add u as chd of v
		}
	};

	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