Submission #871103

#	Time	Username	Problem	Language	Result	Execution time	Memory
871103		MinaRagy06	Paths (RMI21_paths)	C++17	100 / 100	358 ms	24600 KiB

Main

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace std;
using namespace __gnu_pbds;
#define ll long long

const int N = 100'005;
vector<array<int, 2>> adj[N];
tree<pair<ll, int>, null_type, greater<pair<ll, int>>, rb_tree_tag, tree_order_statistics_node_update> s;
int n, k;
ll vals[N];
pair<ll, int> dp[N][2];
void dfs(int i, int par, int parc) {
	dp[i][0] = dp[i][1] = {-1e18, 0};
	for (auto [nxt, w] : adj[i]) {
		if (nxt == par) continue;
		dfs(nxt, i, w);
		pair<ll, int> nxtval = dp[nxt][0];
		nxtval.first += w;
		if (nxtval > dp[i][1]) {
			dp[i][1] = nxtval;
		}
		if (dp[i][0] < dp[i][1]) {
			swap(dp[i][0], dp[i][1]);
		}
	}
	if (adj[i].size() == 1) {
		dp[i][0] = {0, i};
	}
	vals[dp[i][0].second] += parc;
}
ll ans[N];
ll sum = 0;
bool rem(int i) {
	int pos = s.order_of_key({vals[i], i});
	s.erase({vals[i], i});
	if (pos < k) {
		sum -= vals[i];
		return 1;
	}
	return 0;
}
void add(int i, bool gud) {
	s.insert({vals[i], i});
	int pos = s.order_of_key({vals[i], i});
	if (gud) {
		if (pos < k) {
			sum += vals[i];
		} else {
			sum += (*s.find_by_order(k - 1)).first;
		}
	} else {
		if (pos < k) {
			sum += vals[i];
			if (k < s.size()) {
				sum -= (*s.find_by_order(k)).first;
			}
		}
	}
}
void dfs2(int i, int par, int parc, pair<ll, int> dpup) {
	dpup.first += parc;
	bool gud = rem(dp[i][0].second);
	vals[dp[i][0].second] -= parc;
	add(dp[i][0].second, gud);
	ans[i] = sum;
	for (auto [nxt, w] : adj[i]) {
		if (nxt == par) continue;
		pair<ll, int> mx = dp[i][0];
		if (dp[i][0].second == dp[nxt][0].second) mx = dp[i][1];
		mx = max(mx, dpup);
		gud = rem(mx.second);
		vals[mx.second] += w;
		add(mx.second, gud);
		dfs2(nxt, i, w, mx);
		gud = rem(mx.second);
		vals[mx.second] -= w;
		add(mx.second, gud);
	}
	gud = rem(dp[i][0].second);
	vals[dp[i][0].second] += parc;
	add(dp[i][0].second, gud);
}
int main() {
	ios_base::sync_with_stdio(0), cin.tie(0);
	cin >> n >> k;
	for (int i = 1, u, v, w; i < n; i++) {
		cin >> u >> v >> w;
		adj[u].push_back({v, w});
		adj[v].push_back({u, w});
	}
	dfs(1, 0, 0);
	for (int i = 1; i <= n; i++) {
		add(i, 0);
	}
	dfs2(1, 0, 0, {0, 0});
	for (int i = 1; i <= n; i++) {
		cout << ans[i] << '\n';
	}
	return 0;
}

Compilation message (stderr)

Main.cpp: In function 'void add(int, bool)':
Main.cpp:56:10: warning: comparison of integer expressions of different signedness: 'int' and '__gnu_pbds::detail::bin_search_tree_set<std::pair<long long int, int>, __gnu_pbds::null_type, std::greater<std::pair<long long int, int> >, __gnu_pbds::detail::tree_traits<std::pair<long long int, int>, __gnu_pbds::null_type, std::greater<std::pair<long long int, int> >, __gnu_pbds::tree_order_statistics_node_update, __gnu_pbds::rb_tree_tag, std::allocator<char> >, std::allocator<char> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   56 |    if (k < s.size()) {
      |        ~~^~~~~~~~~~

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