The request to minimize the difference mentioned in the problem’s statement is meant to ‘hide’ from the contestant the fact that this difference is either 0 or 1, for any test case. If N is divisible by 3, then the difference is 0; otherwise, it is 1.

There are several correct solutions to the problem, using either DP or greedy algorithms. The author’s solution uses such a greedy algorithm. Let’s consider the sub-tree rooted at vertex R. The algorithm will assign a color (company) to every vertex (employee), in such a way that the difference is minimized. The algorithm works in the following way:

if R is a leaf in the tree, then R is assigned color 1;

if R has only one son (i.e. it is the supervisor of only one employee), then the sub-tree rooted at its son is colored first; then, R is assigned the color which is used the least number of times in the coloring of its son’s sub-tree and which is different from the color assigned to its son; because of the way the algorithm works, the sub-tree rooted at R will be correctly colored (that is, such a color always exists);

if R has 2 sons, LS (the left son) and RS (the right son), the sub-trees rooted at LS and RS are colored independently; then, the colors of the sub-tree rooted at RS are reassigned in such a way that the coloring of the 2 sub-trees minimizes the specified difference, the colors assigned to LS and RS are different and the third color (the one which will be assigned to R) is one of the colors which is used the least number of times in the combined coloring of the 2 sub-trees; the reassignment of the colors in the sub-tree rooted at RS consists in finding a permutation of the colors (that is, every color will be replaced by its corresponding color in the permutation); obviously, by permuting the colors, all the conditions in the sub-tree still hold.

By using the above-mentioned algorithm, the sub-tree rooted at R will be colored correctly. If the algorithm is used for vertex #1, then a solution for the problem is found. The most important part of the algorithm is the third case (when R has 2 sons). Finding a permutation of the colors in the right sub-tree is done by trying each of the 6 possible permutations. Such a permutation always exists; this can be proven by analyzing all the potential cases for the colorings of the left and right sub-trees (there are less then 10 cases).

In order to assign the final color to a vertex V, we keep the permutation of the colors in the sub-tree rooted at V (if V is not the right son of some vertex, then this permutation will be [1 2 3]). The tree is traversed for a second time. For every vertex reached during this traversal, the actual permutation of the colors is already known. When moving towards the son of some vertex, the actual permutation and the permutation corresponding to the son are combined (pretty much like multiplying 2 permutations) and the actual permutation for the next vertex is obtained.

Solution:

#include <stdio.h>

#include <stdlib.h>

#include <string.h>

#define filein "comp.in"

#define fileout "comp.out"

#define maxn 16001

#define mare 1000000

long nf[maxn],fiu[maxn][2],final[maxn],col[maxn],howmany[maxn][3],perm[maxn][3];

long psol[3],pact[3],i,j,k,n,p,hm,c,nmax,nmin;

void df(long node)

{

long fs,fd;

for (k=0;k<3;k++)

perm[node][k]=k;

if (nf[node]==0)

{

col[node]=0;

howmany[node][0]=1;

}

else

if (nf[node]==1)

{

fs=fiu[node][0];

df(fs);

for (k=0;k<3;k++)

howmany[node][k]=howmany[fs][k];

hm=mare; c=0;

for (k=0;k<3;k++)

if (howmany[node][k]<hm && col[fs]!=k)

{

hm=howmany[node][k];

c=k;

}

howmany[node][c]++;

col[node]=c;

}

else /* nf[node]==2 */

{

fs=fiu[node][0];

fd=fiu[node][1];

df(fs);

df(fd);

/* combine the 2 colorings */

for (i=0;i<3;i++)

psol[i]=i;

for (i=0;i<3;i++)

{

pact[0]=i;

for (j=0;j<3;j++)

if (j!=i)

{

pact[1]=j;

for (k=0;k<3;k++)

if (k!=i && k!=j)

{

pact[2]=k;

for (p=0;p<3;p++)

howmany[node][p]=howmany[fs][p];

for (p=0;p<3;p++)

howmany[node][pact[p]]+=howmany[fd][p];

hm=mare;

for (p=0;p<3;p++)

if (howmany[node][p]<hm)

hm=howmany[node][p];

c=0;

for (p=0;p<3;p++)

if (howmany[node][p]>c)

c=howmany[node][p];

nmax=0;

nmin=0;

for (p=0;p<3;p++)

if (howmany[node][p]==hm)

nmin++;

else

nmax++;

if (c-hm<=1 &&

(nmin==3 ||

(nmin==2 && (howmany[node][col[fs]]!=hm ||

howmany[node][pact[col[fd]]]!=hm)) ||

(nmin==1 && (howmany[node][col[fs]]!=hm &&

howmany[node][pact[col[fd]]]!=hm))

) &&

(pact[col[fd]]!=col[fs])

)

{

/* OK! */

for (p=0;p<3;p++)

psol[p]=pact[p];

}

}

}

}

for (k=0;k<3;k++)

perm[fd][k]=psol[k];

for (k=0;k<3;k++)

howmany[node][k]=howmany[fs][k];

for (k=0;k<3;k++)

howmany[node][perm[fd][k]]+=howmany[fd][k];

/* assign a color to node */

hm=mare; c=0;

for (k=0;k<3;k++)

if (howmany[node][k]<hm && col[fs]!=k && perm[fd][col[fd]]!=k)

{

hm=howmany[node][k];

c=k;

}

howmany[node][c]++;

col[node]=c;

}

}

void df2(long node)

{

long p[3];

/* save the current permutation */

for (k=0;k<3;k++)

p[k]=pact[k];

/* modify the permutation */

for (k=0;k<3;k++)

psol[k]=pact[perm[node][k]];

for (k=0;k<3;k++)

pact[k]=psol[k];

final[node]=pact[col[node]];

if (nf[node]>1)

df2(fiu[node][1]);

if (nf[node]>0)

df2(fiu[node][0]);

/* restore the original permutation */

for (k=0;k<3;k++)

pact[k]=p[k];

}

int main()

{

freopen(filein,"r",stdin);

freopen(fileout,"w",stdout);

scanf("%ld",&n);

memset(nf,0,sizeof(nf));

for (k=1;k<n;k++)

{

scanf("%ld %ld",&i,&j);

nf[j]++;

fiu[j][nf[j]-1]=i;

}

memset(col,0,sizeof(col));

memset(howmany,0,sizeof(howmany));

memset(perm,0,sizeof(perm));

df(1);

memset(final,0,sizeof(final));

for (k=0;k<3;k++)

pact[k]=k;

df2(1);

for (i=1;i<=n;i++)

printf("%ld ",final[i]+1);

printf("n");

return 0;

}