Official

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

Claude Opus 5 (High)

Overview

This problem asks us to find the number of ways, modulo \(998244353\), to replace each C in a string \(S\) consisting of A, B, and C with A or B such that the resulting string contains no contiguous substring AA. This can be solved with a standard linear DP (similar to digit DP) that maintains the state of the immediately preceding character.

Analysis

Naive Approach

Let \(k\) be the number of occurrences of C. There are \(2^k\) possible replacements in total. Since \(k\) can be up to \(2 \times 10^6\), an exhaustive search is obviously impossible (\(2^{2000000}\) ways). Therefore, we need an efficient counting method.

Key Observation

The forbidden condition is “no AA appears”, which is a condition that depends only on two adjacent characters. In other words, when deciding the character for day \(i\), the only information we need to know is

whether the character on day \(i-1\) was A or B.

What the characters before that were has no impact whatsoever on whether the choices from day \(i\) onwards are valid.

Because of this property that “we only need information about the single immediately preceding character”, we can manage two DP values:

  • \(\mathrm{dpA}[i]\): the number of valid assignments up to day \(i\) such that the character on day \(i\) is A
  • \(\mathrm{dpB}[i]\): the number of valid assignments up to day \(i\) such that the character on day \(i\) is B

By doing this, we can find the total count in \(O(N)\) time. By grouping states that share the same “last character”, we reduce the time complexity from exponential to linear.

DP Transitions

If we place A on day \(i\), day \(i-1\) must be B (placing A would form AA). Therefore:

\[\mathrm{dpA}[i] = \mathrm{dpB}[i-1]\]

If we place B on day \(i\), day \(i-1\) can be either A or B. Therefore:

\[\mathrm{dpB}[i] = \mathrm{dpA}[i-1] + \mathrm{dpB}[i-1]\]

However, if \(S_i\) is A, we cannot place B on day \(i\) (\(\mathrm{dpB}[i]=0\)), and if \(S_i\) is B, we cannot place A (\(\mathrm{dpA}[i]=0\)). Only when \(S_i\) is C do we apply both transitions.

Concrete Example

For \(S = \) ACA:

  • \(i=1\): Since \(S_1 = \) A, \(\mathrm{dpA}=1, \mathrm{dpB}=0\).
  • \(i=2\): Since \(S_2 = \) C, we try both:
    • Placing A \(\to \mathrm{dpA}[2] = \mathrm{dpB}[1] = 0\) (invalid because it forms AA)
    • Placing B \(\to \mathrm{dpB}[2] = 1 + 0 = 1\)
  • \(i=3\): Since \(S_3 = \) A, only A is possible:
    • \(\mathrm{dpA}[3] = \mathrm{dpB}[2] = 1\), \(\mathrm{dpB}[3] = 0\)

The answer is \(1 + 0 = 1\) (only ABA).

Note that if \(S\) originally contains AA (e.g., \(S=\) AA), both values will become \(0\) during the transitions, and the answer will automatically become \(0\). No special case handling is needed.

Algorithm

  1. Process day 1 to set up the initial values:
    • If \(S_1 = \) A \(\to (\mathrm{dpA}, \mathrm{dpB}) = (1, 0)\)
    • If \(S_1 = \) B \(\to (0, 1)\)
    • If \(S_1 = \) C \(\to (1, 1)\)
  2. For \(i = 2, 3, \dots, N\), compute the following (let the new values be \(nA, nB\), initialized to \(0\)):
    • When \(S_i \in \{\)A, C\(\}\): \(nA \gets \mathrm{dpB}\)
    • When \(S_i \in \{\)B, C\(\}\): \(nB \gets (\mathrm{dpA} + \mathrm{dpB}) \bmod 998244353\)
    • \((\mathrm{dpA}, \mathrm{dpB}) \gets (nA, nB)\)
  3. The answer is \((\mathrm{dpA} + \mathrm{dpB}) \bmod 998244353\).

Complexity

  • Time complexity: \(O(N)\)
  • Space complexity: \(O(N)\) (for storing the input string; the DP itself uses \(O(1)\))

Implementation Notes

  • Use two variables instead of maintaining an array: Since \(\mathrm{dp}[i]\) depends only on \(\mathrm{dp}[i-1]\), an array of length \(N\) is unnecessary. While using an array will pass since \(N \le 2\times 10^6\), using just two variables is friendlier to both memory and constant factor performance.

  • Be careful with the update order: If you overwrite \(\mathrm{dpA}\) first, the old value of \(\mathrm{dpA}\) cannot be used in the calculation of \(\mathrm{dpB}\). Always store the new values in temporary variables \(nA, nB\) before assigning them simultaneously.

  • Take modulo % at each step: \(\mathrm{dpA}+\mathrm{dpB}\) is at most around \(2 \times (\text{MOD}-1)\), so it will not overflow a long long (or even an int if handled properly), but make sure to take modulo to keep the values in \([0, \text{MOD})\).

  • Fast I/O: Since \(N\) is large (\(2\times 10^6\)), use scanf / std::cin with sync_with_stdio(false) in C++, or read from sys.stdin in Python. Reading the string character-by-character with slow methods risks getting TLE.

  • For the \(N = 1\) case, simply adding the initial values and outputting the result works correctly (for C, the answer is \(2\)).

    Source Code

#include <bits/stdc++.h>
using namespace std;
const long long MOD = 998244353;
int main(){
    int n;
    scanf("%d", &n);
    static char buf[2000006];
    scanf("%s", buf);
    string s(buf);
    // dpA: ways ending with A, dpB: ways ending with B
    long long dpA = 0, dpB = 0;
    // first char
    if(s[0]=='A'){dpA=1;dpB=0;}
    else if(s[0]=='B'){dpA=0;dpB=1;}
    else {dpA=1;dpB=1;}
    for(int i=1;i<n;i++){
        long long nA=0,nB=0;
        char c=s[i];
        if(c=='A'||c=='C') nA = dpB;
        if(c=='B'||c=='C') nB = (dpA+dpB)%MOD;
        dpA=nA;dpB=nB;
    }
    printf("%lld\n", (dpA+dpB)%MOD);
    return 0;
}

This editorial was generated by claude5opus-high.

posted:
last update: