Submission #1056906

#	Time	Username	Problem	Language	Result	Execution time	Memory
1056906		jamesbamber	Split the Attractions (IOI19_split)	C++17	40 / 100	512 ms	1048576 KiB

split

#include "split.h"
#include <iostream>

using namespace std;

vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
	vector<int> ids = {1, 2, 3};
	if(a > b) swap(a, b), swap(ids[0], ids[1]);
	if(b > c) swap(b, c), swap(ids[1], ids[2]);
	if(a > b) swap(a, b), swap(ids[0], ids[1]);

	vector<vector<int>> ad(n);
	for(int i=0; i<p.size(); i++){
		if(n == p.size() and i == 0) continue;
		ad[p[i]].push_back(q[i]);
		ad[q[i]].push_back(p[i]);
	}


	vector<int> vis(n);
	auto dfs = [&](auto &&dfs, int v, int p, vector<int> &cc) -> void {
		cc.push_back(v);
		vis[v] = 1;
		for(int u: ad[v]){
			if(vis[u] or u == p) continue;
			dfs(dfs, u, v, cc);
		}
	};


	pair<int,int> split = {-1, -1};
	vector<int> sz(n);
	auto calc_sz = [&](auto &&calc_sz, int v, int p) -> void {
		sz[v] = 1;
		for(int u: ad[v]){
			if(u == p) continue;
			calc_sz(calc_sz, u, v);
			sz[v] += sz[u];
		}	
		if(sz[v] >= a and n - sz[v] >= b) split = {v, p};
		if(sz[v] >= b and n - sz[v] >= a) split = {p, v};
	};

	vector<int> ans(n, 0);

	if(a == 1){
		vector<int> order;
		dfs(dfs, 0, -1, order);
		for(int i=0; i<b; i++) ans[order[i]] = ids[1];
		for(int i=0; i<n; i++){
			if(ans[i] != 0) continue;
			if(a) ans[i] = ids[0], a--;
			else if(c) ans[i] = ids[2], c--;
		}
	}
	else{
		calc_sz(calc_sz, 0, -1);
		if(split.first == -1) return ans;
		vector<int> cca, ccb;
		dfs(dfs, split.first, split.second, cca);
		dfs(dfs, split.second, split.first, ccb);

		for(int i=0; i<a; i++) ans[cca[i]] = ids[0];
		for(int i=0; i<b; i++) ans[ccb[i]] = ids[1];
		for(int i=0; i<n; i++){
			if(ans[i] == 0) ans[i] = ids[2];
		}
	}

	return ans;
}

Compilation message (stderr)

split.cpp: In function 'std::vector<int> find_split(int, int, int, int, std::vector<int>, std::vector<int>)':
split.cpp:13:16: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   13 |  for(int i=0; i<p.size(); i++){
      |               ~^~~~~~~~~
split.cpp:14:8: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   14 |   if(n == p.size() and i == 0) continue;
      |      ~~^~~~~~~~~~~

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