G - At Most 2 Colors Editorial by spheniscine


One could try to define \(dp_i\) as the number of valid strings of length \(i\) and do DP on it, but it isn’t quite obvious how to define the transitions.

Instead, let’s add the constraint that \(dp_i\) is the number of valid strings of length \(i\) containing at least two colors. With this definition the following formula can be derived for \(i \ge 2\):

\[ dp_i = 2\cdot dp_{i-1} + (C-2)dp_{i+1-K} + C(C-1)\]

(note that \(dp_i = 0\) for \(i \le 1\), including for negative \(i\))

Explanation:

\(2\cdot dp_{i-1}\) : If the color at position \(i\) is equal to one of the two most-recently used colors, we can take a valid string of length \(i-1\) and append that color to it.

\( (C-2)dp_{i+1-K}\) : If the color at position \(i\) isn’t one of the two most-recently used colors, the last \(K-1\) positions must be of a single color. So we can take a valid string of length \(i-1-(K-2) = i+1-K\), and repeat the last color of that string \(K-2\) more times to create a block of at least \(K-1\). We then can append one of \(C-2\) colors to it.

\(C(C-1)\) : There are \(C(C-1)\) new valid strings of the form AA…AAB.

The final answer is \(dp_N + C\), as we have \(C\) valid strings of exactly one color that wasn’t counted in the DP.

posted:
last update: