D - シフトの割り当て / Shift Assignment 解説 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:
When \(S_i = \)
A: The \(i\)-th character must beA. If the previous character wasA, it formsAA, which is invalid; this is only allowed when the previous character wasB(or at the start of the string). $\(a_i = b_{i-1},\qquad b_i = 0\)$When \(S_i = \)
B: The \(i\)-th character must beB. The previous character can be anything. $\(a_i = 0,\qquad b_i = a_{i-1} + b_{i-1}\)$When \(S_i = \)
C: We sum the cases for choosingAand choosingB(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) % Mallows 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
0and terminate. This handles cases whereAAis 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 withfor ch in syields integer values for each element (Ais 65,Bis 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.
投稿日時:
最終更新: