Submission #559774

#	Time	Username	Problem	Language	Result	Execution time	Memory
559774		flappybird	Road Closures (APIO21_roads)	C++17	100 / 100	196 ms	34628 KiB

roads

#include "roads.h"

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
#define MAX 101010

int N;

vector<pii> adj[MAX];
int w[MAX];

typedef multiset<ll, greater<ll>> msl;
msl maxSet[MAX], st[MAX];
ll X[MAX];
int prt[MAX];
int K;
int deg[MAX];
ll sum;

void getW(int x, int p = -1) {
	prt[x] = p;
	for (auto& [v, c] : adj[x]) {
		if (v == p) continue;
		X[v] = w[v] = c;
		maxSet[x].insert(X[v]);
		sum += X[v];
		getW(v, x);
	}
}

bool exist(msl& ms, ll v) {
	return ms.find(v) != ms.end();
}

void refresh(int v) {
	while (st[v].size() && (*st[v].begin()) > *maxSet[v].rbegin()) {
		int x = *st[v].begin();
		int y = *maxSet[v].rbegin();
		sum -= y;
		sum += x;
		st[v].erase(st[v].find(x));
		maxSet[v].erase(maxSet[v].find(y));
		maxSet[v].insert(x);
		st[v].insert(y);
	}
}

void changev(int v, ll prv, ll nxt) {
	if (exist(maxSet[v], prv)) {
		maxSet[v].erase(maxSet[v].find(prv));
		maxSet[v].insert(nxt);
		sum -= prv;
		sum += nxt;
	}
	else {
		st[v].erase(st[v].find(prv));
		st[v].insert(nxt);
	}
	refresh(v);
}

ll getK(int v) {
	if (maxSet[v].size() < K) return 0;
	return *maxSet[v].rbegin();
}

void calc(int v) {
	ll prevX = X[v];
	refresh(v);
	while (maxSet[v].size() > K) {
		st[v].insert(*maxSet[v].rbegin());
		int x = *maxSet[v].rbegin();
		maxSet[v].erase(maxSet[v].find(x));
		sum -= x;
	}
	X[v] = max(w[v] - getK(v), 0ll);
	if (X[v] == prevX) return;
	if (!v) return;
	changev(prt[v], prevX, X[v]);
	calc(prt[v]);
}

vector<ll> minimum_closure_costs(int N, vector<int> U, vector<int> V, vector<int> W) {
	::N = N;
	int i;
	ll sumW = 0;
	for (i = 0; i < N - 1; i++) adj[U[i]].emplace_back(V[i], W[i]), adj[V[i]].emplace_back(U[i], W[i]), sumW += W[i], deg[U[i]]++, deg[V[i]]++;
	getW(0);
	vector<int> degSort;
	for (i = 0; i < N; i++) degSort.push_back(i);
	sort(degSort.begin(), degSort.end(), [&](int i, int j) { return deg[i] > deg[j]; });
	vector<ll> ans(N);
	ans[N - 1] = sum;
	vector<int> changeList;
	int ptr = 0;
	for (K = N - 2; K >= 1; K--) {
		while (ptr < N && K < deg[degSort[ptr]]) changeList.push_back(degSort[ptr++]);
		for (auto v : changeList) calc(v);
		ans[K] = sum;
	}
	for (auto& v : ans) v = sumW - v;
	return ans;
}

Compilation message (stderr)

roads.cpp: In function 'll getK(int)':
roads.cpp:65:23: warning: comparison of integer expressions of different signedness: 'std::multiset<long long int, std::greater<long long int> >::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
   65 |  if (maxSet[v].size() < K) return 0;
      |      ~~~~~~~~~~~~~~~~~^~~
roads.cpp: In function 'void calc(int)':
roads.cpp:72:26: warning: comparison of integer expressions of different signedness: 'std::multiset<long long int, std::greater<long long int> >::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
   72 |  while (maxSet[v].size() > K) {
      |         ~~~~~~~~~~~~~~~~~^~~

• Subtask 1: 5 / 5
• Subtask 2: 7 / 7
• Subtask 3: 14 / 14
• Subtask 4: 10 / 10
• Subtask 5: 17 / 17
• Subtask 6: 25 / 25
• Subtask 7: 22 / 22