This problem is equivalent to the following graph theory formulation: given a tree ~T~ with ~N~ nodes, color the nodes with colors ~1, 2, 3, \dots~, such that no two adjacent nodes have the same color, and that the sum of all colors is minimal. Since a tree is a bipartite graph, it can be always be colored with two colors ~1~ and ~2~, but this does not always yield the minimal solution as shown in the following simple example:

First, we need to construct the adjacency lists for each node. This can be achieved by using vectors or lists in order to use ~\mathcal O(N)~ memory.

Assume that the number of different colors used in the minimal solution is ~C~. We can formulate a recursive solution (though note that this can also be computed bottom-up by removing leaves from the tree as we go).

For each node ~v~, let ~f(v, c)~ be the minimal sum of colors in the subtree rooted at node ~v~, given that ~c \in \{1, 2, \dots, C\}~ is the color of node ~v~. Let ~v_1, v_2, \dots, v_k~ be the children of ~v~. We can try any of the colors ~1, 2, \dots, C~ for the node ~v~, and compute ~f(v, c)~ as:

$$\displaystyle f(v, c) = c + \sum_{i=1}^k \min_{c_i \ne c} f(v_i, c_i)$$

Therefore, direct implementation of the above method gives us a solution with time and memory complexity ~\mathcal O(NC^2)~.

We can also prove that ~\lfloor \log_2 N \rfloor~ is an upper bound for ~C~. Let ~C(k)~ be the size of the smallest tree that needs all colors ~1, 2, \dots, k~ in an optimal coloring. Trivially, it holds that ~C(1) = 1~ and ~C(2) = 2~. Without loss of generality, we can pick the node with color ~k~ to be the root. In that case, the root needs to be adjacent to all colors from ~1~ to ~k-1~ and we can apply the inductive hypothesis as follows:

$$\displaystyle C(k) \ge C(k-1) + \dots + C(2) + C(1) + 1 \ge 2^{k-1} + \dots + 2^1 + 2^0 + 1 = 2^k$$

This completes the proof. The above algorithm therefore has complexity ~\mathcal O(N \log^2 N)~ in both time and memory. We leave the problem of constructing a minimal tree that requires ~k~ colors as an exercise for the reader.

However, there is also an algorithm with ~\mathcal O(N)~ time and memory complexity. The above formula can be simplified if we just use the two best values for each node, together with the colors where these minimum values are achieved.

Let ~c_i~ and ~d_i~ be the best and second-best colors for the node ~v_i~ respectively. Now,

$$\displaystyle f(v, c) = c + \sum_{i=1}^k \begin{cases} f(v_i, c_i) & \text{if }c \ne c_i \\ f(v_i, d_i) & \text{otherwise} \end{cases}$$

This appears to still take ~\mathcal O(Ck)~ time as we need to try all ~C~ colors, leading to an ~\mathcal O(N \log N)~ solution. However, we can improve this to ~\mathcal O(k)~ time. Let ~B~ be our base cost, the minimum we must pay for the subtree rooted at ~v~:

$$\displaystyle B = c + \sum_{i=1}^k f(v_i, c_i)$$

Now, let ~A_i~ be the additional cost if we decide to use color ~i~:

$$\displaystyle A_i = \sum_j [f(v_j, d_j) - f(v_j, c_j)\text{, where }c_j = i]$$

We can precompute ~A~ in ~\mathcal O(k)~ time, and then we can get the minimum cost of coloring ~v~'s subtree in ~\mathcal O(k)~ time with:

$$\displaystyle B + \min(A_1, \dots, A_k)$$

And with that, our solution is now ~\mathcal O(N)~!