Submission #509943

#	Time	Username	Problem	Language	Result	Execution time	Memory
509943		daongocha	Paths (RMI21_paths)	C++14	100 / 100	562 ms	31504 KiB

Main

#pragma GCC optimize("Ofast")
#pragma GCC optimize("unroll-loops")
#pragma GCC target("avx,avx2,fma")
 
#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace __gnu_pbds;
 
// #define map dsfakuhwkl34hdfs
#define int long long
 
#define fileopen(a, b) freopen(((std::string)a + ".inp").c_str(), "r", stdin); freopen(((std::string)b + ".out").c_str(), "w", stdout);
#define fileopen1(a) freopen(((std::string)a + ".inp").c_str(), "r", stdin); freopen(((std::string)a + ".out").c_str(), "w", stdout);
 
using namespace std;
 
const int MAXN = 3e5 + 5;
 
#define ordered_set tree<pair<int, int>, null_type,less<pair<int, int>>, rb_tree_tag,tree_order_statistics_node_update>
 
vector<pair<int, int>> adj[MAXN];
pair<int, int> st[MAXN << 1];
int n, k;
int dep[MAXN];
int tin[MAXN], tout[MAXN];
int ver[MAXN];
 
void build() {
	for (int i = 1; i <= n; i++) {
		st[tin[i] + n - 1] = {dep[i], i};
	} 
	for (int i = n - 1; i >= 1; i--) st[i] = max(st[i << 1], st[i << 1 | 1]);
}
 
void upd(int i, int val) {
	i += n - 1;
	for (st[i].first = val; i > 1; i >>= 1)
		st[i >> 1] = max(st[i], st[i ^ 1]);
}
 
pair<int, int> query(int l, int r) {
	l += n - 1;
	r += n;
	pair<int, int> ans = {-1, 0};
	for (; l < r; l >>= 1, r >>= 1) {
		if (l & 1) ans = max(ans, st[l++]);
		if (r & 1) ans = max(ans, st[--r]);
	}
	return ans;
}
 
int par[MAXN];
int num[MAXN];
bool vis[MAXN];
 
int update(int u) {
	int ans = 0;
	for (; !vis[u]; u = par[u]) {
		ans += num[u];
		vis[u] = 1;
	}
	return ans;
}
 
 
void dfs(int u, int pre = 0) {
	static int timer = 0;
 
	tin[u] = ++timer;
 
	par[u] = pre;
	for (auto vw: adj[u]) {
		int v = vw.first, w = vw.second;
		if (v != pre) {
			dep[v] = dep[u] + w;
			num[v] = w;
			dfs(v, u);
		}
	}
 
	tout[u] = timer;
}
 
ordered_set pp;
int ans = 0;
 
stack<pair<int, int>> exist;
 
void add(int u, int w, bool sv) {
	int v = pp.order_of_key({-dep[u], u});
	pp.erase({-dep[u], u}); 
	auto it = *(pp.find_by_order(k - 1));
	if (v < k) {
		ans -= dep[u];
		ans -= it.first;
	}
	
	dep[u] += w;
	pp.insert({-dep[u], u});
	
	if (it.first > -dep[u]) {
		ans += dep[u];
		ans += it.first;
	}
 
	// cout << tin[u] << ' ' << dep[u] << " owo\n";
	upd(tin[u], dep[u]);
}
 
void reroot_down(int u, int v, int w) {
	// maximal in subtree v and out of subtree v
	// for (int i = 1; i <= n; i++) cout << dep[tin[i]] << ' '; cout << '\n';
 
	auto it1 = query(tin[v], tout[v]); 
	auto it2 = max(query(1, tin[v] - 1), query(tout[v] + 1, n));
 
	add(it1.second, -w, 1);
	add(it2.second, w, 1);
}
 
void reroot_up(int u, int v, int w) {
	swap(u, v);
	auto it1 = query(tin[v], tout[v]);
	auto it2 = max(query(1, tin[v] - 1), query(tout[v] + 1, n));
 
	// cout << v << '\n';
	// cout << 1 << ' ' << tin[v] - 1 << ' ' << tout[v] + 1 << ' ' << n << '\n';
 
	add(it1.second, w, 1);
	add(it2.second, -w, 1);
}
 
 
int res[MAXN];
 
void dfs2(int u, int pre = 0) {
	for (auto vw: adj[u]) {
		int v = vw.first, w = vw.second;
		if (v != pre) {
			reroot_down(u, v, w);
			res[v] = ans;
			dfs2(v, u);
			// cout << v << '\n';
			reroot_up(v, u, w);
		}
	}
}
 
signed main() {
	#ifndef PICHU_LOCAL_DEF
		//fileopen1("LAH");
	#endif
 
	cin.tie(0)->sync_with_stdio(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});
	} 
 
	vis[0] = 1;
 
	dfs(1);
 
	#define pq dsaflkdfssdfdfs
	vector<pair<int, int>> pq;
	for (int i = 1; i <= n; i++) 
		pq.push_back({dep[i], i});
 
	sort(pq.rbegin(), pq.rend());
 
	for (int i = 0; i < n; i++) {
		int u = pq[i].second;
		dep[u] = update(u);
	}
 
	#undef pq
 
	build();
 
	for (int i = 1; i <= n; i++) pp.insert({-dep[i], i});	
 
	pp.insert({100, 0});
	
	for (int i = 0; i < k; i++) ans += pp.find_by_order(i)->first;
	ans = -ans;
 
	res[1] = ans;
		
	// for (int i = 1; i <= n; i++) cout << tin[i] << ' '; cout << '\n';
 
	dfs2(1);
	
	// for (int i = 1; i <= n; i++) cout << dep[i] << ' '; cout << '\n';
	// reroot_down(1, 5, 1);
	// for (int i = 1; i <= n; i++) cout << dep[i] << ' '; cout << '\n';
	// // reroot_up(5, 1, 1);
	// for (int i = 1; i <= n; i++) cout << dep[i] << ' '; cout << '\n';
 
	for (int i = 1; i <= n; i++) cout << res[i] << '\n';
}	

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