1879E - Interactive Game with Coloring - CodeForces Solution

#include<bits/stdc++.h>

using namespace std;

const int N = 123;

int n;
int color[N];
int countColors[N][N];
int p[N];
vector<int> g[N];
int deg[N];

void add_edge(int x, int y)
{
	g[x].push_back(y);
	g[y].push_back(x);
}

bool tryTwoColors()
{
	int v1 = n + 1;
	int v2 = n + 2;
	for(int i = 2; i <= n; i++)
	{
		if(p[i] != 1)
		{
			add_edge(i, p[i]);
		}
	}
	for(int i = 2; i <= n; i++)
	{
		if(deg[i] == 1)
			add_edge(i, v1);
	}
	for(int i = 2; i <= n; i++)
	{
		if(p[i] != 1 && deg[p[i]] == 1)
			add_edge(i, v2);
	}
	add_edge(v1, v2);
	bool bad = false;
	for(int i = 2; i <= n + 2; i++)
		if(color[i] == 0)
		{
			color[i] = 1;
			queue<int> q;
			q.push(i);
			while(!q.empty())
			{
				int k = q.front();
				q.pop();
				for(auto y : g[k])
				{
					if(color[y] == 0)
					{
						color[y] = 3 - color[k];
						q.push(y);
					}
					else if(color[y] == color[k])
						bad = true;
				}
			}
		}
	if(bad)
		for(int i = 2; i <= n + 2; i++)
			color[i] = 0;
	return !bad;
}

void tryThreeColors()
{
	for(int i = 2; i <= n; i++)
		if(p[i] == 1)
			color[i] = 1;
		else
			color[i] = (color[p[i]] % 3) + 1;
}

int findVertex(const vector<int>& colors)
{
	int s = colors.size();
	for(int i = 2; i <= n; i++)
	{
		if(vector<int>(countColors[i], countColors[i] + s) == colors)
			return i;
	}
	return -1;
}

int main()
{
	cin >> n;
	for(int i = 2; i <= n; i++)
	{
		cin >> p[i];
		deg[p[i]]++;
	}
	
	if(*max_element(p + 2, p + n + 1) == 1)
	{
		for(int i = 2; i <= n; i++)
			color[i] = 1;
	}
	else if (!tryTwoColors())
		tryThreeColors();
	
	int colorsUsed = *max_element(color + 2, color + n + 1);
	
	cout << colorsUsed << endl;
	for(int i = 2; i <= n; i++)
	{
		cout << color[i];
		if(i == n) cout << endl;
		else cout << " ";
	}
	cout.flush();
	
	for(int i = 2; i <= n; i++)
	{
		countColors[i][color[i]]++;
		countColors[p[i]][color[i]]++;
	}
	
	while(true)
	{
		int resp;
		cin >> resp;
		if(resp == -1 || resp == 1)
			exit(0);
		vector<int> counts(colorsUsed + 1);
		for(int i = 1; i <= colorsUsed; i++)
			cin >> counts[i];
		int v = findVertex(counts);
		assert(v != -1);
		cout << color[v] << endl;
		cout.flush();
	}
}