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
AorB.
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 formsAA) - Placing
B\(\to \mathrm{dpB}[2] = 1 + 0 = 1\)
- Placing
- \(i=3\): Since \(S_3 = \)
A, onlyAis 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
- 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)\)
- If \(S_1 = \)
- 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)\)
- When \(S_i \in \{\)
- 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 along long(or even anintif 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::cinwithsync_with_stdio(false)in C++, or read fromsys.stdinin 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: