Official

D - シフトの割り当て / Shift Assignment Editorial by admin

gpt-6-astra-medium

Overview

We replace undecided shifts C with A or B and find the number of assignments where AA does not appear. We process characters from left to right and use dynamic programming, maintaining the number of ways based on whether the trailing character is A or B.

Analysis

If there are \(K\) occurrences of C, there are \(2^K\) possible ways to replace them in total. Trying all possibilities will not run in time since \(N\) is up to \(2 \times 10^6\).

Here, let us consider appending a new character to the end of a valid string.

  • A can only be appended if the preceding character is B.
  • B can be appended regardless of whether the preceding character is A or B.

In other words, to determine whether the next character can be appended, we only need the last character, not the entire string up to that point.

Therefore, for the processed prefix, we categorize and count the valid assignments into the following \(2\) types:

  • \(a\): the number of assignments ending with A
  • \(b\): the number of assignments ending with B

By updating these two values as we proceed, we do not need to enumerate all possible replacements.

Algorithm

1. Initialization

Before processing any characters, we set:

\(a = 0,\quad b = 1\)

Although an empty string has no last character, for convenience, we include the single state where either A or B can be placed at the beginning into \(b\). You can think of it as having a hypothetical B immediately preceding the string.

2. Update from Left to Right

Letting the values after the update be \(a'\) and \(b'\), the transitions based on the current character are as follows:

Current character \(a'\) \(b'\)
A \(b\) \(0\)
B \(0\) \(a+b\)
C \(b\) \(a+b\)

The reasons for each are as follows:

When the current character is A

It is fixed that A is placed. To avoid AA, we only carry over assignments where the preceding character is B.

\(a' = b,\quad b' = 0\)

When the current character is B

It is fixed that B is placed. Since either character is allowed immediately before, we sum the number of ways for both cases.

\(a' = 0,\quad b' = a+b\)

When the current character is C

Both A and B can be chosen.

  • If we choose A, we only carry over assignments where the preceding character is B.
  • If we choose B, we carry over all assignments.

\(a' = b,\quad b' = a+b\)

Take the result modulo \(998244353\) after each addition.

3. Finding the Answer

After processing all characters, we sum the cases where the string ends with A and B.

The answer is \((a+b) \bmod 998244353\).

At each transition, we only perform additions that do not create AA. Furthermore, every valid assignment is counted in exactly one transition path, so the total count is found without any omissions or duplicates.

Concrete Example

Consider the case where \(S = \texttt{CCA}\).

Processed prefix \(a\) \(b\)
Empty string \(0\) \(1\)
C \(1\) \(1\)
CC \(1\) \(2\)
CCA \(2\) \(0\)

There are \(2\) valid strings, ABA and BBA, so the answer is \(2\).

Complexity

  • Time complexity: \(O(N)\)
  • Space complexity: \(O(N)\), required to store the input string. The additional space for DP is \(O(1)\) since only a few variables are used.

Implementation Notes

  • Use the values before updating.
    For example, when processing C, simultaneous assignment like a, b = b, total allows setting the new \(a\) to the previous value of \(b\).

  • The modulo can be computed with a single subtraction.
    Since we always maintain \(0 \le a,b < \mathrm{MOD}\), we have \(a+b < 2\mathrm{MOD}\). Therefore, if the sum is greater than or equal to \(\mathrm{MOD}\), subtracting \(\mathrm{MOD}\) once is sufficient.

  • The input string is processed as bytes.
    When iterating over a byte string read with sys.stdin.buffer.readline, each element is an integer. Therefore, in the code, A is compared with 65 and B with 66.

  • No special handling is needed for impossible cases.
    For example, if there is a fixed AA, the transition results in \(a=b=0\). Both remain \(0\) from then on, naturally outputting \(0\) as the answer.

    Source Code

import sys

def main():
    input = sys.stdin.buffer.readline
    n = int(input())
    s = input().strip()

    MOD = 998244353
    a, b = 0, 1

    for c in s:
        if c == 65:  # A
            a, b = b, 0
        else:
            total = a + b
            if total >= MOD:
                total -= MOD
            if c == 66:  # B
                a, b = 0, total
            else:  # C
                a, b = b, total

    answer = a + b
    if answer >= MOD:
        answer -= MOD
    print(answer)

if __name__ == "__main__":
    main()

This editorial was generated by gpt-6-astra-medium.

posted:
last update: