Submission #1150925

#TimeUsernameProblemLanguageResultExecution timeMemory
1150925gygSplit the Attractions (IOI19_split)C++20
22 / 100
551 ms1114112 KiB
#include "split.h"
#include <bits/stdc++.h>
using namespace std;
#define sig signed
#define int long long
#define arr array 
#define vec vector
#define pii pair<int, int>
#define fir first 
#define sec second
const int N = 1e5 + 5;

int n;
arr<int, 4> nm;
arr<vec<int>, N> adj;

arr<int, N> sz;
void sz_dfs(int u, int pr = -1) {
	sz[u] = 1;
	for (int v : adj[u]) {
		if (v == pr) continue;
		sz_dfs(v, u);
		sz[u] += sz[v];
	}
}
int cnt_dfs(int u, int pr = -1) {
	for (int v : adj[u]) {
		if (v == pr) continue;
		if (2 * sz[v] > n) return cnt_dfs(v, u);
	}
	return u;
}
arr<set<int>, N> sb;
void sb_dfs(int u, int src, int pr) {
	sb[src].insert(u);
	for (int v : adj[u])
		if (v != pr) sb_dfs(v, src, u);
}

arr<int, N> ans;
void mrk_dfs(int u, set<int>& x, int y, int pr = -1) {
	if (!nm[y]) return;
	ans[u] = y, nm[y]--;
	for (int v : adj[u]) {
		if (v == pr) continue;
		if (!x.count(v)) continue;
		mrk_dfs(v, x, y, u);
	}
}

vec<sig> find_split(sig _n, sig _a, sig _b, sig _c, vec<sig> _u, vec<sig> _v) {
	n = _n, nm[1] = _a, nm[2] = _b, nm[3] = _c;
	for (int i = 0; i < _u.size(); i++) {
		int u = _u[i] + 1, v = _v[i] + 1;
		adj[u].push_back(v), adj[v].push_back(u);
	}
	vec<pii> srt = {{nm[1], 1}, {nm[2], 2}, {nm[3], 3}};
	sort(srt.begin(), srt.end());
	nm[1] = srt[0].fir, nm[2] = srt[1].fir, nm[3] = srt[2].fir;

	sz_dfs(1);
	int cnt = cnt_dfs(1);
	for (int u : adj[cnt]) sb_dfs(u, u, cnt);

	// cout << cnt << '\n';
	// for (int u : adj[cnt]) cout << u << ": " << sb[u].size() << '\n';

	bool flg = false;
	for (int u : adj[cnt]) {
		if (sb[u].size() < nm[1]) continue;
		flg = true;
		set<int> rm;
		for (int v = 1; v <= n; v++) rm.insert(v);
		for (int v : sb[u]) rm.erase(v);

		// for (int v : sb[u]) cout << v << " ";
		// cout << '\n';
		// for (int v : rm) cout << v << " ";
		// cout << '\n';

		mrk_dfs(*sb[u].begin(), sb[u], 1);
		mrk_dfs(*rm.begin(), rm, 2);
		for (int v = 1; v <= n; v++)
			if (!ans[v]) ans[v] = 3, nm[3]--;
		break;
	}
	if (!flg) { vec<sig> rsp(n, 0); return rsp; }

	vec<sig> rsp;
	for (int u = 1; u <= n; u++) rsp.push_back(srt[ans[u] - 1].sec);
	return rsp;
}
#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...