This is my solution to coloring the graph (white & red) which I feel is simpler than the one given in the official editorial; but it boils down to the same thing.
This is only for k=2. k=1 needs to be handled separately and is fairly straightforward.
Solution:
-
Precompute the sizes of all subtrees
in the tree rooted at 1. -
Color the root vertex white.
-
Now do this recursively, for every
vertex v; starting from root vertex -
For every child c of v, if size of c
is even, simply color c opposite of
v. -
Now for odd children of v, color half
the children (n/2) same as v and half
the children (n+1/2) opposite of v.
That’s it. How does this ensure all the constraints are satisfied?
Correctness:
Constraint - connectedness
- As for the constraint that H_2 should be fully connected with only edges with opposite colored ends, if is enough to prove every edge (u, v) in the original tree has either
- opposite colored ends OR 2) there exists a third node w such that (w, u) and (w, v) are both edges with opposite colored ends. (This will make sure there is a path u -> w -> v from u to v satisfying the condition)
-
Consider an edge (u, v) where v is the child of u. If u and v are opposite colored, we’re done.
-
Else if u and v are same colored, we know that u must have another child w which is of opposite color (because of the half-half coloring of odd children) making (w, u) and (w, v) opposite colored.
Constraint - equal white and red nodes
-
Lemma: Every even sized subtree has equal number of white and red nodes. Similarly, every odd sized subtree rooted at u satisfies |W - R| = 1 and the color which is in excess is precisely the color of root u.
-
Proof: By induction,
Base Case: Leaf node, it is odd sized (size 1) and it has excess of its own color
Consider subtree rooted at node u. Without loss of generality, assume it is white (W)
Case - size(u) is even
Split its children subtrees into even and odd. For every even child, there are equal number of W and R nodes. Now for odd children, it is easy to prove that there will be odd number of odd children. (n-1/2) of the children have excess of W and (n+1/2) children have excess of R. Therefore, total W nodes (n-1/2 + root_node) = total R nodes (n+1/2). Hence, for even sized subtree, W = R
Case - size(u) is odd
Again, split its children subtrees into even and odd. For every even child, equal number of W and R nodes. For odd children, you can prove that there will be even number of odd children. (n/2) of the children will have excess W and (n/2) of the children will have excess R, thus balancing the number of W and R in all the children. But, for the full subtree there will be an excess of W due to the root node. Hence, for odd sized subtree, |W - R| = 1 and that excess color is exactly the color of the respective root node.
Hope it helps! Suggestions are welcome 
Here is my code :