Submission #143521

#	Time	Username	Problem	Language	Result	Execution time	Memory
143521		mosesmayer	Split the Attractions (IOI19_split)	C++17	100 / 100	212 ms	20176 KiB

split

#include "split.h"
#include <bits/stdc++.h>
#define fi first
#define se second
#define eb emplace_back
#define pb push_back
#define ALL(x) x.begin(), x.end()
using namespace std;
typedef pair<int,int> pii;
typedef vector<int> vi;

struct Dsu{
	int n;
	vi par;
	Dsu() {}
	Dsu(int n): n(n), par(n) {}
	void init(){
		iota(par.begin(), par.end(), 0);
	}
	int fpr(int u){
		return par[u] == u ? u : par[u] = fpr(par[u]);
	}
	bool merge(int u, int v){
		if ((u = fpr(u)) == (v = fpr(v))) return 0;
		par[v] = u;
		return 1;
	}
};

const int mxsz = 1e5 + 5;
int n, m;
Dsu dsu;
vector<pii> ord, edges;
vi g[mxsz], adj[mxsz];
vi ans;
int sz[mxsz];
int bfspar[mxsz], bfsord[mxsz];

void buildSpanningTree(int st){
	fill(bfspar, bfspar + n, -1);
	queue<int> q;
	q.push(st);
	bfspar[st] = st;
	int cntr = 0;
	while (!q.empty()){
		int u = q.front(); q.pop();
		bfsord[cntr++] = u; sz[u] = 1;
		for (int nx : g[u]){
			if (bfspar[nx] == -1){
				bfspar[nx] = u;
				adj[u].pb(nx);
				adj[nx].pb(u);
				q.push(nx);
			}
		}
	}
	bfspar[st] = -1;
	for (int i = n-1; i > 0; --i)
		sz[bfspar[bfsord[i]]] += sz[bfsord[i]];
}
int findCentroid(){
	int cpar = -1, u = 0;
	while (cpar != u){
		bool work = 0;
		for (int nx : adj[u]){
			if (nx == cpar) continue;
			if (sz[nx] > n / 2){
				cpar = u; u = nx; work = 1; break;
			}
		}
		if (!work) cpar = u;
	}
	return u;
}
void get_sz(int u, int par){
	sz[u] = 1;
	for (int nx : adj[u]){
		if (nx == par) continue;
		get_sz(nx, u);
		sz[u] += sz[nx];
	}
}
vi res;
queue<int> q;
void workA(int st, int prv){
//	puts("Run1");
	q.push(st);
	for (int i = 0; i < ord[0].fi; i++){
		int u = q.front(); q.pop();
		if (res[u] != 0) {
			i--; continue;
		}
//		printf("%d ", u);
		res[u] = ord[0].se;
		for (int nx : adj[u]){
			if (nx != prv && res[nx] == 0){
				q.push(nx);
			}
		}
	}
	while (!q.empty()) q.pop();
//	puts("\nRun2");
	q.push(prv);
	for (int i = 0; i < ord[1].fi; i++){
		int u = q.front(); q.pop();
		if (res[u] != 0) {
			i--; continue;
		}
//		printf("%d ", u);
		res[u] = ord[1].se;
		for (int nx : adj[u]){
			if (nx != st && res[nx] == 0){
				q.push(nx);
			}
		}
	}
//	puts("\nDone");
	for (int i = 0; i < n; i++)
		if (res[i] == 0) res[i] = ord[2].se;
}
void work(int x, int cen){
//	printf("Work %d cen %d\n", x, cen);
	q.push(x);
//	puts("W1");
	for (int i = 0; i < ord[0].fi; i++){
		int u = q.front(); q.pop();
		if (res[u] != 0) {
			i--; continue;
		}
//		printf("%d ", u);
		res[u] = ord[0].se;
		for (int nx : g[u]){
			if (res[nx] == 0 && dsu.fpr(nx) == x){
				q.push(nx);
			}
		}
	}
	while (!q.empty()) q.pop();
//	puts("\nW2");
	q.push(cen);
	for (int i = 0; i < ord[1].fi; i++){
		int u = q.front(); q.pop();
		if (res[u] != 0) {
			i--; continue;
		}
//		printf("%d ", u);
		res[u] = ord[1].se;
		for (int nx : g[u]){
			if (res[nx] == 0){
				q.push(nx);
			}
		}
	}
//	puts("\nDone");
	for (int i = 0; i < n; i++)
		if (res[i] == 0) res[i] = ord[2].se;
}

void compressRegion(int u, int head, int prv){
	dsu.merge(head, u);
//	printf("%d %d %d\n", u, head, prv);
	for (int nx : adj[u]){
		if (nx != prv) compressRegion(nx, head, u);
	}
}

void Main(){
//	for (int i = 0; i < n; i++){
//		printf("%d: ", i);
//		for (int nx : g[i]){
//			printf(" %d", nx);
//		}
//		puts("");
//	}
	buildSpanningTree(0);
//	puts("Spanning tree:");
//	for (int i = 0; i < n; i++){
//		printf("%d: ", i);
//		for (int nx : adj[i]){
//			printf(" %d", nx);
//		}
//		puts("");
//	}
	int cen = findCentroid();
//	printf("Centroid: %d\n", cen);
	get_sz(cen, -1);
	assert(sz[cen] == n);
	res.resize(n); fill(ALL(res), 0);
	for (int nx : adj[cen]){
		if (sz[nx] >= ord[0].fi){ // has subtree >= a
//			printf("Case 1: %d\n", nx);
			workA(nx, cen);
			return;
		}
	}
	for (int i = 0; i < adj[cen].size(); i++){
		compressRegion(adj[cen][i], adj[cen][i], cen);
	}
//	for (int i = 0; i < n; i++){
//		printf("%d par %d\n", i, dsu.fpr(i));
//	}
	for (int i = 0, x, y; i < m; i++){
		tie(x, y) = edges[i];
		x = dsu.fpr(x); y = dsu.fpr(y);
		if (x == cen || y == cen) continue;
		if (dsu.merge(x, y)){
			sz[x] += sz[y];
			if (sz[x] >= ord[0].fi){
				work(x, cen);
				return;
			}
		}
	}
}
vi find_split(int N, int A, int B, int C, vi P, vi Q){
	n = N, m = P.size();
	dsu = Dsu(N);
	dsu.init();
	ord.eb(A, 1); ord.eb(B, 2); ord.eb(C, 3);
	sort(ALL(ord));
//	for (int i = 0; i < 3; i++){
//		printf("(%d, %d)%c", ord[i].fi, ord[i].se, " \n"[i==2]);
//	}
	for (int i = 0; i < m; i++){
		g[P[i]].pb(Q[i]);
		g[Q[i]].pb(P[i]);
		edges.eb(P[i], Q[i]);
	}
	Main();
	return res;
}

Compilation message (stderr)

split.cpp: In function 'void Main()':
split.cpp:196:20: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
  for (int i = 0; i < adj[cen].size(); i++){
                  ~~^~~~~~~~~~~~~~~~~

• Subtask 1: 7 / 7
• Subtask 2: 11 / 11
• Subtask 3: 22 / 22
• Subtask 4: 24 / 24
• Subtask 5: 36 / 36