Submission #390292

#	Time	Username	Problem	Language	Result	Execution time	Memory
390292		JimmyZJX	Split the Attractions (IOI19_split)	C++14	22 / 100	106 ms	15144 KiB

split

#include <iostream>
#include <fstream>
#include <algorithm>
#include <cstring>
#include <climits>
#include <cassert>
#include <tuple>
#include <queue>
#include <stack>
#include <vector>
#include <map>
#include <set>
#include <string>
#include <unordered_set>
#include <unordered_map>

using namespace std;

typedef long long LL;
typedef vector<int> Vi;
typedef vector<bool> Vb;
typedef vector<vector<int>> Vii;

#define forR(i, n) for (int i = 0; i < (n); i++)

Vii nbs(100003);
Vii children(100003);
Vi result;

void task_tree(int n, Vi abc) {
	Vb vis(n);
	Vi visVec;
	queue<int> q; q.push(0);

	while (!q.empty()) {
		int x = q.front(); q.pop();
		visVec.push_back(x); vis[x] = true;

		for (int nb : nbs[x]) {
			if (vis[nb]) continue;
			children[x].push_back(nb);
			q.push(nb);
		}
	}

	int treeNode, nodeColor, rootColor;
	 
	Vi sz(n);
	for (auto it = visVec.rbegin(); it != visVec.rend(); it++) {
		int x = *it;
		int sumCnt = 0;
		for (int nb : children[x]) sumCnt += sz[nb];
		sz[x] = sumCnt + 1;

		forR(i, 3) forR(j, 3) {
			if (i == j) continue;
			if (sz[x] >= abc[i] && n - sz[x] >= abc[j]) {
				treeNode = x; nodeColor = i + 1; rootColor = j + 1;
				goto FoundSolution;
			}
		}
	}
	return;

FoundSolution:
	// 1: color node
	int nodeCnt = abc[nodeColor - 1];
	q = queue<int>(); q.push(treeNode);
	for (int i = 0; i < nodeCnt; i++) {
		int x = q.front(); q.pop();
		result[x] = nodeColor;
		for (int child : children[x]) q.push(child);
	}

	// 2. color root
	int rootCnt = abc[rootColor - 1];
	q = queue<int>(); q.push(0);
	for (int i = 0; i < rootCnt; i++) {
		int x = q.front(); q.pop();
		result[x] = rootColor;
		for (int child : children[x])
			if (result[child] == 0)
				q.push(child);
	}

	// 3. color rest
	int colorRest = 6 - nodeColor - rootColor;
	for (int& c : result) {
		if (c == 0) c = colorRest;
	}
}

Vi find_split(int n, int a, int b, int c, Vi p, Vi q) {
	int m = p.size();
	result = Vi(n);

	forR(i, m) {
		nbs[p[i]].push_back(q[i]);
		nbs[q[i]].push_back(p[i]);
	}

	if (m == n - 1) {
		// Task 3, Tree
		task_tree(n, Vi{ a, b, c });
	}

	return result;
}


#ifdef TEST_LOCAL
int main() {
	find_split(6, 2, 2, 2, Vi{ 0, 1, 0, 0, 0 }, Vi{ 1, 2, 3, 4, 5 });

	return 0;
}
#endif

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