Editorial for An Animal Contest 6 P5 - Rock Painting

Remember to use this editorial only when stuck, and not to copy-paste code from it. Please be respectful to the problem author and editorialist.
Submitting an official solution before solving the problem yourself is a bannable offence.

Author: rickyqin005

Subtask 1 Hint

What is the probability of rock i being the first or last rock with some colour c?

Subtask 1

Let's create an array a of length N to represent the initial state of the rocks. If a_i = 0, the i^\text{th} rock has not been coloured. Otherwise, a_i will denote the colour of the i^\text{th} rock.

We also define f(i) to be the probability of rock i being the rightmost rock and g(i) to be the probability of rock i being the leftmost rock with some colour c. We have:

\displaystyle f(i) = \begin{cases}
\frac{(C-k_i)(C-1)^{x_i}}{C^{x_i+1}} & \text{if }a_i = 0 \\
0 & \text{if }a_i \ne 0\text{ and there exists some }j > i\text{ with }a_i = a_j \\
(\frac{C-1}{C})^{x_i} & \text{otherwise}

where x_i is the number of non-coloured rocks to the right of rock i and k_i is the number of distinct non-zero a_i to the right of rock i. Calculating g(i) is the same as f(i) but in reverse.

Therefore, the final answer to this problem is:

\displaystyle \sum_{i=1}^N i \cdot f(i) - \sum_{i=1}^N i \cdot g(i)

Time Complexity: \mathcal{O}(N + \log(10^9 + 7))

Subtask 2

For this subtask, we can generalize the definition of f(i) to be f(i)=(\frac{C-1}{C})^{N-i}. Using r = \frac{C-1}{C}, we have:

\displaystyle \begin{align}
S_{f(i)} &= \sum_{i=1}^N i \cdot r^{N-i} \\
&= 1 \cdot r^{N-1} + 2 \cdot r^{N-2} + \dots + (N-1) \cdot r^1 + N \cdot r^0

Using some algebra, we have:

\displaystyle \begin{align}
rS_{f(i)} &= 1 \cdot r^N + 2 \cdot r^{N-1} + \dots + (N-1) \cdot r^2 + N \cdot r^1 \\
rS_{f(i)} - S_{f(i)} &= r^N + r^{N-1} + \dots + r^2 + r^1 - N \cdot r^0 \\
S_{f(i)} &= \frac{r^1 + r^2 + \dots + r^{N-1} + r^N - N}{r-1}

Substituting the sum of geometric series formula into the equation, we get:

\displaystyle \begin{align}
S_{f(i)} &= \frac{\frac{r(r^N-1)}{r-1} - N}{r-1} \\
&= \frac{r(r^N-1)}{(r-1)^2} - \frac{N}{r-1}

Time Complexity: \mathcal{O}(\log(10^9 + 7))

Subtask 3

This subtask is an extension of subtasks 1 and 2 and is left as an exercise for the reader.

Time Complexity: \mathcal{O}(M \log(10^9 + 7))


There are no comments at the moment.