Official

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

Claude Opus 5 (High)

Overview

We are given a string \(S\) consisting of A, B, and C. The problem asks us to find the number of ways, modulo \(998244353\), to replace each C with either A or B such that the resulting string does not contain AA. This is a typical counting problem with a “no consecutive identical elements” condition, which can be solved using digit DP (or character-by-character DP).

Analysis

Naive Approach

Let \(k\) be the number of Cs; there are \(2^k\) ways to replace them. Since \(N \le 2\times 10^6\), \(k\) can be up to \(2\times 10^6\), making complete enumeration impossible (\(2^{2000000}\) possibilities).

Key Observation

The condition of not containing AA is a local condition: “no two adjacent characters are both A”. That is, when determining the \(i\)-th character, whether a violation occurs depends solely on whether the \((i-1)\)-th character was A.

Therefore, when determining characters sequentially from left to right, the only information we need to remember is:

  • The string up to this point satisfies the condition
  • Whether the immediately preceding character was A

These are the only two states needed. By maintaining the number of valid configurations for each state and accumulating them, we can reduce exponential enumeration to linear-time DP.

Algorithm

Consider the following DP:

  • \(a_i\): The number of ways to determine the first \(i\) characters such that the condition is satisfied and the \(i\)-th character is A
  • \(b_i\): The number of ways to determine the first \(i\) characters such that the condition is satisfied and the \(i\)-th character is B

For the base case, considering that “the \(0\)-th character is not A”, we set \(a_0 = 0,\ b_0 = 1\) (it is natural to interpret \(b\) as the state where “the previous character is not A”).

The transitions based on the \(i\)-th character \(S_i\) are as follows:

  1. When \(S_i = \) A: The \(i\)-th character must be A. If the previous character was A, it forms AA, which is invalid; this is only allowed when the previous character was B (or at the start of the string). $\(a_i = b_{i-1},\qquad b_i = 0\)$

  2. When \(S_i = \) B: The \(i\)-th character must be B. The previous character can be anything. $\(a_i = 0,\qquad b_i = a_{i-1} + b_{i-1}\)$

  3. When \(S_i = \) C: We sum the cases for choosing A and choosing B (combining the two cases above). $\(a_i = b_{i-1},\qquad b_i = a_{i-1} + b_{i-1}\)$

The answer is \(a_N + b_N \bmod 998244353\). If no valid string can be formed, it automatically becomes \(0\) (e.g., if \(S = \) AA, then \(a_1=1, b_1=0 \Rightarrow a_2 = b_1 = 0, b_2 = 0\)).

Concrete Example

For \(S = \) CAC:

\(i\) Character \(a_i\) \(b_i\)
0 - 0 1
1 C \(b_0=1\) \(a_0+b_0=1\)
2 A \(b_1=1\) 0
3 C \(b_2=0\) \(a_2+b_2=1\)

The answer is \(0 + 1 = 1\). In fact, only BAB satisfies the condition (AAB, AAA, and BAA all contain AA).

Complexity

  • Time Complexity: \(O(N)\)
  • Space Complexity: \(O(N)\) (to store the input string; the DP itself is \(O(1)\))

Implementation Details

  • Do not allocate an array: Since \(a_i\) and \(b_i\) depend only on the immediately preceding values, we can reuse two variables. Since \(N = 2\times 10^6\), allocating an array of length \(N\) wastes both memory and time. In Python, using tuple assignment like a, b = b, (a+b) % M allows updating without temporary variables.

  • Take the modulo at each step: Take the remainder modulo \(998244353\) at each step to prevent \(a+b\) from growing excessively large (in Python, integers have arbitrary precision, so failing to do so will cause the number of digits to explode and run extremely slowly).

  • Early termination: Once \(a = b = 0\), all subsequent transitions will remain \(0\), so you can immediately print 0 and terminate. This handles cases where AA is guaranteed to appear very quickly.

  • Fast I/O: Because \(N\) is large, read the entire input at once using sys.stdin.buffer.read(). Iterating directly over a byte sequence with for ch in s yields integer values for each element (A is 65, B is 66), which allows for faster comparisons than string character comparisons.

  • Even with \(N \ge 1\), to guard against input format variations where \(S\) might be read as empty, it is safer to include a check like len(data) > 1.

    Source Code

import sys

def main():
    data = sys.stdin.buffer.read().split()
    n = int(data[0])
    s = data[1] if len(data) > 1 else b""
    M = 998244353
    a = 0  # ends with A
    b = 1  # ends with B (or start)
    A = 65
    B = 66
    for ch in s:
        if ch == A:
            a, b = b, 0
        elif ch == B:
            a, b = 0, (a + b) % M
        else:
            a, b = b, (a + b) % M
        if b == 0 and a == 0:
            print(0)
            return
    print((a + b) % M)

main()

This editorial was generated by claude5opus-high.

posted:
last update: