Official

F - Wine Thief Editorial by evima


Let us find the sequence \(b\) where \(b_i\) is the number of ways to steal bottles, including the \(i\)-th bottle from the left, without getting noticed by Takahashi.
If we find it, the answer is \(\displaystyle \sum_{i = 1}^N A_ib_i\).

It would be easy if \(b_1 = b_2 = b_3 = \dots = b_N\). However, because the row of bottles has ends, the bottles have not much symmetry, besides the horizontal one.
For now, let us consider a slightly different version of the problem where \(N\) bottles are arranged in a circle, that is, Bottles \(1\) and \(N\) are also adjacent.
In this case, all \(N\) bottles are symmetric to each other, so \(b_1 = b_2 = b_3 = \dots = b_N = \frac{xK}{N}\), where \(x\) is the number of all ways to steal bottles without getting noticed.

Let \(f(n, k)\) be the number of ways to steal \(k\) bottles from \(n\) bottles arranged in a circle, where no two stolen bottles are adjacent.
Also, let \(g(n, k)\) be the number of ways to steal \(k\) bottles from \(n\) arranged in a line, where no two stolen bottles are adjacent.
We can see that \(f(n, k)\) is equal to \(g(n - 3, k - 1) + g(n - 1, k)\), by dividing such ways into the ones where Bottle \(1\) is stolen and the ones where it is not stolen.
Additionally, \(g(n, k)\) is equal to the number of ways to arrange \(k\) black stones and \(n - k\) white stones so that no two black stones are adjacent. We can represent this number as \(_{k + 1} \mathrm{H}_{n - k - (k - 1)}\) using combination with repetition, by starting from a alternating row with \(k\) black stones and \(k - 1\) white stones and adding \(n - k - (k - 1)\) white stones to it.

Now, we have solved the circular version of the problem. If the bottles are arranged in a line, we can divide the ways to steal bottles into the following two categories:

  • The ways that are also allowed in the circular version: As we discussed above, each of \(b_1, b_2, b_3, \dots, b_N\) gets added by \(\frac{f(N, K)K}{N}\).

  • The ways that choose both Bottles \(1\) and \(N\), which is not allowed in the circular version: Such ways contribute \(g(N - 4, K - 2)\) to each of \(b_1\) and \(b_N\).
    The contributions to \(b_3, b_4, b_5, \dots, b_{N - 2}\) can be recursively found by letting \(N \leftarrow N - 4, K \leftarrow K - 2\).

If we do pre-computation to enable ourselves to compute \(_n\mathrm{C}_r\) in \(O(1)\) time, we can find \(f(n, r)\) and \(g(n, r)\) in \(O(1)\) time.
Since there are \(O(N)\) recursions, if we efficiently do additions to ranges in \(b\) by, for example, using accumulated sums, we can compute \(b\) in \(O(N)\) time, allowing us to find the final answer.

Watch out the corner cases such as \(N = K = 1\) that can appear during the recursive process.

posted:
last update: