Solving " Tricolor Arrangement " Efficiently

A rectangular board with three rows and $n$ columns is filled with $3n$ counters, of which $n$ are red, $n$ are white, and $n$ are blue. The object is to rearrange the counters to have counters of each of the three different colors in every column. The only operation allowed is to swap counters in the same row. Design an algorithm to accomplish such a task or prove that such an algorithm does not exist.

Let $\mathcal{O}(1)$ be the cost of swapping the counters in the same row.

Iterate over each column. Let us consider that up to column $i-1$, we are done. Now we are at column $i$, there will be three counters, if they are of different colors then we are done, else fix one counter ( say white ) and scan other two rows to get two different color counters ( red and blue ). Running the previous step for a constant number of times ( six ) will solve the problem. More precisely there will be six possibilities in each column (123,132,.....,321), so I will try each of those.

The runtime of the above algorithm will be $\mathcal{O}(n^3)$ as there are $n$ columns and time spent on each column is $\mathcal{O}(n^2)$.

Question: Is it possible to solve the above-described problem linear time $\mathcal{O}(n)$ time as input here has size $\mathcal{O}(n)$? I am not even able to come up with an algorithm faster than $\mathcal{O}(n^3)$ time.

I have taken this question from the book titled Algorithmic Puzzles ( problem 46 )

Here is an algorithm with $O(n)$ time-complexity.

First, count how many counters of each color there are on each row. With $O(n)$ time, we will obtain the number of red counters $r_i$, the number of white counters $w_i$, the number of blue counters $b_i$ on row $i$, $i=1,2,3$.

Second, we can solve the following equations for nonnegative integer variables $a_k$, $1\le k\le6$ in $O(1)$ time.

$$r_1 = a_1 + a_2$$
$$w_1 = a_3 + a_4$$
$$b_1 = a_5 + a_6$$
$$r_2 = a_4 + a_5$$
$$w_2 = a_1 + a_6$$
$$b_2 = a_2 + a_3$$
$$r_3 = a_3 + a_6$$
$$w_3 = a_2 + a_5$$
$$b_3 = a_1 + a_4$$

Third, we arrange the three rows into the following configuration.
$$
\begin{array}{|l|l|l|l|l|l|} \hline
a_1\text{ red} & a_2\text{ red} & a_3\text{ white} & a_4\text{ white} & a_5\text{ blue} & a_6\text{ blue} \\\hline
a_1\text{ white} & a_2\text{ blue} & a_3\text{ blue} & a_4\text{ red} & a_5\text{ red} & a_6\text{ white} \\\hline
a_1\text{ blue} & a_2\text{ white} & a_3\text{ red} & a_4\text{ blue} & a_5\text{ white} & a_6\text{ red} \\\hline
\end{array}
$$

We are done.

Here are a couple of exercises that fill the gap in the above specification of the algorithm.

Exercise 1. Find an algorithm that solves in $O(1)$ time that system of 9 equations for nonnegative integer variables $a_k$, $1\le k\le6$, given the condition that $r_i+w_i+b_i=n$ for all $i$ and $\sum_{i=1}^3r_i=\sum_{i=1}^3w_i=\sum_{i=1}^3b_i=n$. Show that there is always a solution.

Exercise 2. Check that arrangement in third step can be done in $O(n)$ time by swapping the counters in each row.

Here is a problem that may not be easy to do.

Problem 1. Generalize the problem of tricolor arrangement to $m$ rows and $m$ colors, $m>3$. Solve it with an $O(n)$ algorithm, assuming $m$ is a constant. (Hint, use Hall's marriage theorem to reduce the sum of the number of distinct colors in each row.)

StackExchange Computer Science Q#102165, answer score: 2