Submission #258625

#TimeUsernameProblemLanguageResultExecution timeMemory
258625NightlightSplit the Attractions (IOI19_split)C++14
100 / 100
153 ms20204 KiB
#include "split.h"
#include <bits/stdc++.h>
#define pii pair<int, int>
#define pib pair<int, bool>
#define mp make_pair
using namespace std;
 
vector<int> adj[200005];
vector<int> ans;
int N, rt, M;
pii id[5];
int cnt[5];
int sz[100005], par[100005];
int ch;
int col[100005], vis[100005];
int parent[100005];
bool found;
 
void DFS(int u) {
	sz[u] = 1;
	vis[u] = 1;
	for(int v : adj[u]) {
		if(vis[v]) continue;
		DFS(v);
		sz[u] += sz[v];
    parent[v] = u;

	}
}
 
void cent() {
	DFS(0);
	int u = 0, now = N / 2, nxt;
	bool ok = 0;
	while(!ok) {
		ok = 1, nxt = u;
		for(int v : adj[u]) {
			if(sz[v] > now && sz[v] < sz[u]) {
				nxt = v, ok = 0;
				break;
			}
		}
		u = nxt;
	}
	rt = u;
}
 
int findpar(int u) {
	return par[u] == u ? u : par[u] = findpar(par[u]);
}
 
void DFS1(int u) {
	if(id[1].first <= 0) return;
	vis[u] = 2;
	id[1].first--;
	ans[u] = id[1].second;
	for(int v : adj[u]) {
		if(vis[v] == 2 || col[v] == 0) continue;
		DFS1(v);
	}
}
 
void DFS2(int u) {
	if(id[2].first <= 0) return;
	vis[u] = 2;
	id[2].first--;
	ans[u] = id[2].second;
	for(int v : adj[u]) {
		if(vis[v] == 2) continue;
		DFS2(v);
	}
}
 
vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
	N = n, id[1] = {a, 1}, id[2] = {b, 2}, id[3] = {c, 3}, M = p.size();
	sort(id + 1, id + 4);
	ans.resize(N, 0);
	for(int i = 0; i < M; i++) {
		adj[p[i]].emplace_back(q[i]);
		adj[q[i]].emplace_back(p[i]);
	}
	cent();
	for(int i = 0; i < N; i++) par[i] = i, sz[i] = 1;
	int u, v;
	bool ok = 0;
//	cout << rt << "\n";
	for(int i = 0; i < M; i++) {
		if(p[i] == rt || q[i] == rt) continue;
    if(!(parent[p[i]] == q[i] || parent[q[i]] == p[i])) continue;
		u = findpar(p[i]), v = findpar(q[i]);
		if(sz[u] < sz[v]) swap(u, v);
		if(u == v) continue;
		par[v] = u;
		sz[u] += sz[v];
	}
	for(int i = 0; i < M; i++) {
		if(p[i] == rt || q[i] == rt) continue;
		u = findpar(p[i]), v = findpar(q[i]);
		if(sz[u] < sz[v]) swap(u, v);
		if(sz[u] >= id[1].first) continue;
		if(u == v) continue;
		par[v] = u;
		sz[u] += sz[v];
	}
	for(int i = 0; i < N; i++) {
		if(i == rt) continue;
		u = findpar(i);
		if(sz[u] >= id[1].first) {
//			cout << u << "\n";
			ok = 1;
			for(int j = 0; j < N; j++) {
				if(findpar(j) == u) col[j] = 1;
			}
			break;
		}
	}
	if(ok == 0) return ans;
//	cout << id[1].first << " " << id[2].first << "\n";
	DFS1(u);
	DFS2(rt);
	for(int i = 0; i < N; i++) if(ans[i] == 0) ans[i] = id[3].second;
	return ans;
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...