Official

Ex - Dye Color Editorial by en_translator


For a sequence \(A=(A_1,A_2,\dots,A_N)\), let us denote by \(f(A)\) the expected value of the number of operations. Also, if all elements of \(A\) are equal, let us call it a terminal state.

Let \(E_{A,A'}\) be the probability that a sequence \(A\) changes to a sequence \(A'\) in one operation.

If \(A\) is a terminal state, then \(f(A)=0\); otherwise, \(f(A) = 1 + \sum E_{A,A'} f(A')\).

Also, let us define a sequence \(B\) by \(B_{A,j} := (\)the number of \(i\) such that \(A_i = j)\).

The problem can be solved if a constant \(C\) and a function \(g(x)\) exists such that \( \sum_{i=1}^{N} g(B_{A,i}) = f(A) + C\).

\(g\) should satisfy \( \sum_{i=1}^{N} g(B_{A,i}) =1+ \sum E_{A,A'} \sum_{i=1}^{N} g(B_{A',i})\) for any sequence \(A\) that is not a terminal state.

If \(g\) satisfies \(g(B_{A,i}) = \frac{1}{N} + \sum E_{A,A'} g(B_{A',i})\) for \(1 \le i \le N\), then the equation above always holds.

Moreover, when \(B_i = j\), the probability that a single operation makes \(B_i = k\) can be written in the form only dependent on \(j\) and \(k\).

Therefore, let \(P_{j,k}\) be the probability when \(B_i=j\) that a single operation makes \(B_i=j\), then the condition that \(g\) should satisfy is that \(g(i) = \frac{1}{N} + \sum_{j=0}^{N} P_{i,j} g(j)\) for \(0 \le i < N\). For \(i=N\), it is already in a terminal state, so the equation above does not need to be satisfied.

Moreover, an operation does not increase \(B_i\) by two or more.

Thus, the condition can be expressed with a simultaneous equation as follows:

\(\begin{pmatrix} P_{0,0}-1 & P_{0,1} & 0 & 0 & \dots & 0 & 0\\ P_{1,0} & P_{1,1}-1 & P_{1,2} & 0 & \dots & 0 & 0\\ \vdots & \vdots & \vdots & \vdots & \ddots & \vdots & \vdots\\ P_{N-1,0} & P_{N-1,1} & P_{N-1,2} & P_{N-1,3} & \dots & P_{N-1,N-1}-1 & P_{N-1,N} \\ \end{pmatrix}\) \(\begin{pmatrix} g(0) \\ g(1) \\ \vdots \\ g(N) \\ \end{pmatrix}\) \(=\) \(\begin{pmatrix} -\frac{1}{N} \\ -\frac{1}{N} \\ \vdots \\ -\frac{1}{N} \\ \end{pmatrix}\)

By letting \(g(0)=0\) and applying the \(i\)-th constraint in the order of \(i=1,2,\dots,N-1\), we can find \(g(i)\) in \(\mathrm{O}(N^2)\) time.

However, we still have issues. How can we find \(P_{i,j}\) fast enough? Also, we have to prove that \(P_{i,i+1} \not \equiv 0 \bmod 998244353\).

We will find \(P_{i,j}\) fast. In other words, when there are \(i\) balls in Color \(X\), we will compute the probability that there will be \(j\) balls after a single operation.

Suppose that \(S\) balls in Color \(X\) were first chosen. Then, the probability that Color \(X\) is chosen when changing the color is \(\displaystyle \frac{\displaystyle \sum_{k=0}^{N-i} \binom{N-i}{k} \times \binom{i}{S} \times \frac{k+S}{N}}{2^N}\). Since it can be decomposed to a term dependent on \(k\), a term dependent on \(S\), a term dependent on \(k+S\), and a constant, so the value above can be found in a total of \(\mathrm{O}(N\log N)\) time for all \(S\). By decomposing \(\frac{k+S}{N}\) to \(\frac{k}{N} + \frac{S}{N}\), it can be computed in a total of \(\mathrm{O}(N)\) time too.

The probability that Color \(X\) is not chosen can be similarly found in a total of \(\mathrm{O}(N\log N)\) time.

Therefore, all \(P_{i,j}\) have been computed in a total of \(\mathrm{O}(N^2\log N)\) time.

Moreover, the probability that \(S=0\) and Color \(X\) is chosen when changing the color of balls is \(P_{i,i+1}\), and its value is \(\displaystyle \frac{\displaystyle \sum_{k=0}^{N-i} \binom{N-i}{k} \times \frac{k}{N}}{2^N}\). This value is transformed to \(\displaystyle \frac{N-i}{N \times 2^{i+1}}\), which is non-zero modulo \(998244353\) for \(0 \le i < N\).

Therefore, the problem has been solved in a total of \(\mathrm{O}(N^2\log N)\) or \(\mathrm{O}(N^2)\) time.

posted:
last update: