Submission #1047757

#	Time	Username	Problem	Language	Result	Execution time	Memory
1047757		MercubytheFirst	Split the Attractions (IOI19_split)	C++14	0 / 100	27 ms	8284 KiB

split

#include <bits/stdc++.h>
#include "split.h"
using namespace std;

void my_assert(bool x) {
	if(!x) {while(true) {} }
}

vector<vector<int> > g;
vector<int> sub;
pair<int, int> sz[3];
int cnt[3];
int ansv = -1, ansp = -1;

int sub_dfs(int v, int p) {
	sub[v] = 1;
	for(int to : g[v]) {
		if(to != p) {
			sub[v] += sub_dfs(to, v);
		}
	}
	return sub[v];
}


void finder_dfs(int v, int p) {
	if(ansv != -1) {
		return;
	}
	const int subtree = sub[v], uptree = sub[0] - sub[v];
	if(subtree >= sz[0].first and uptree >= sz[1].first) {
		ansv = v;
		ansp = p;
		return;
	}

	for(int to : g[v]) {
		if(to != p) {
			finder_dfs(to, v);
			if(ansv != -1) {
				return;
			}
		}
	}
}

void finalize_dfs(vector<int>& res, int v, int p, int group) {
	if(cnt[group] >= sz[group].first) {
		// assert(group != 2);
		group = 2;
	}
	assert(!res[v]);
	assert(0 <= group and group < 3);
	res[v] = sz[group].second;
	cnt[group]++;
	for(int to : g[v]) {
		if(to == p or res[to] != 0) {
			continue;
		}
		finalize_dfs(res, to, v, group);
	}
}

vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
	vector<int> res;
	// if (n == 9) {
	// 	res = {1, 1, 3, 1, 2, 2, 3, 1, 3};
	// } else {
	// 	res = {0, 0, 0, 0, 0, 0};
	// }
	sz[0] = {a, 1};
	sz[1] = {b, 2};
	sz[2] = {c, 3};
	sort(sz, sz+3);
	g.assign(n, vector<int>());
	sub.assign(n, 0);
	int m = p.size();
	assert(m == n - 1);
	for(int i = 0; i < m; ++i) {
		g[p[i]].push_back(q[i]);
		g[q[i]].push_back(p[i]);
	}
	sub_dfs(0, -1);
	finder_dfs(0, -1);
	res.assign(n, 0);
	if(ansv == -1) {
		return res;
	}
	const int ansv_sub = sub[ansv], ansv_up = sub[0] - sub[ansv];
	if(ansv_sub < ansv_up) {
		finalize_dfs(res, ansv, ansp, 0);
		finalize_dfs(res, 0, -1, 1);
	}
	else {
		finalize_dfs(res, ansv, ansp, 1);
		finalize_dfs(res, 0, -1, 0);
	}
	for(int i = 0; i < n; ++i) {
		assert(res[i] >= 0);
		if(res[i] == 0) {
			res[i] = 2;
		}
	}
	return res;
}


/*
6 5
1 3 2
0 1      
0 2
0 3
0 4
0 5
*/

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