Submission #538070

#	Time	Username	Problem	Language	Result	Execution time	Memory
538070		hmm789	Paths (RMI21_paths)	C++14	20 / 100	342 ms	21804 KiB

Main

#include <bits/stdc++.h>
using namespace std;
#define int long long 

int k;
vector<pair<int, int>> adj[100000];
int up[100000], down[100000], ans[100000], num[100000], curans;
int pre[100000], post[100000], cur;
multiset<int> s1, s2;

pair<int, int> dfs(int x, int d, int p) {
	pre[x] = cur++;
	pair<int, int> res = make_pair(d, x);
	for(auto i : adj[x]) {
		if(i.first == p) continue;
		pair<int, int> tmp = dfs(i.first, d+i.second, x);
		if(tmp.first > res.first) res = tmp;
	}
	down[x] = res.second;
	post[x] = cur-1;
	return res;
}

int dist[100000], dist2[100000];
void dfs2(int x, int d, int p, bool b) {
	if(b) dist[x] = d;
	else dist2[x] = d;
	for(auto i : adj[x]) if(i.first != p) dfs2(i.first, d+i.second, x, b);
}

void dfs3(int x, int d, int p, int to) {
	if(x == to) num[x] = d;
	for(auto i : adj[x]) {
		if(i.first == p) continue;
		if(down[i.first] == to) dfs3(i.first, d+i.second, x, to);
		else dfs3(i.first, i.second, x, down[i.first]);
	}
}

void insert(int x) {
	if(s2.empty() || x > *(--s2.end())) {
		s1.insert(x);
		curans += x;
	} else s2.insert(x);
	if(s1.size() > k) {
		s2.insert(*s1.begin());
		curans -= *s1.begin();
		s1.erase(s1.begin());
	}
}

void remove(int x) {
	if(!s2.empty() && x <= *(--s2.end())) {
		s2.erase(s2.find(x));
	}
	else {
		curans -= x;
		assert(s1.find(x) != s1.end());
		s1.erase(s1.find(x));
		if(!s2.empty()) {
			s1.insert(*(--s2.end()));
			curans += *(--s2.end());
			s2.erase(--s2.end());
		}
	}
}

void reroot(int x, int p) {
	ans[x] = curans;
	for(auto i : adj[x]) {
		if(i.first == p) continue;
		remove(num[up[i.first]]); remove(num[down[i.first]]);
		num[up[i.first]] += i.second;
		num[down[i.first]] -= i.second;
		insert(num[up[i.first]]); insert(num[down[i.first]]);
		reroot(i.first, x);
		remove(num[up[i.first]]); remove(num[down[i.first]]);
		num[up[i.first]] -= i.second;
		num[down[i.first]] += i.second;
		insert(num[up[i.first]]); insert(num[down[i.first]]);
	}
}

int32_t main() {
	ios_base::sync_with_stdio(0);
	cin.tie(0); cout.tie(0);
	int n, a, b, w;
	cin >> n >> k;
	for(int i = 0; i < n-1; i++) {
		cin >> a >> b >> w;
		a--; b--;
		adj[a].push_back(make_pair(b, w));
		adj[b].push_back(make_pair(a, w));
	}
	dfs2(0, 0, -1, 1);
	int far = 0, far2 = 0;
	for(int i = 0; i < n; i++) if(dist[i] > dist[far]) far = i;
	dfs2(far, 0, -1, 0);
	for(int i = 0; i < n; i++) if(dist2[i] > dist2[far2]) far2 = i;
	dfs2(far2, 0, -1, 1);
	dfs(far, 0, -1);
	for(int i = 0; i < n; i++) {
		if(pre[far] > pre[i] && post[far] < post[i]) up[i] = far2;
		else if(pre[far2] > pre[i] && post[far2] < post[i]) up[i] = far;
		else if(dist[i] > dist2[i]) up[i] = far2;
		else up[i] = far;
	}
	dfs3(far, 0, -1, down[far]);
	for(int i = 0; i < n; i++) insert(num[i]);
	reroot(far, -1);
	for(int i = 0; i < n; i++) cout << ans[i] << '\n';
}

Compilation message (stderr)

Main.cpp: In function 'void insert(long long int)':
Main.cpp:45:15: warning: comparison of integer expressions of different signedness: 'std::multiset<long long int>::size_type' {aka 'long unsigned int'} and 'long long int' [-Wsign-compare]
   45 |  if(s1.size() > k) {
      |     ~~~~~~~~~~^~~

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