D - シフトの割り当て / Shift Assignment Editorial by admin
Gemini 3.8 Flash (High)Overview
This problem asks us to find the number of ways, modulo \(998244353\), to replace each C in a string \(S\) of length \(N\) with A or B such that there are no two consecutive As.
Analysis
Naive Approach
Letting \(K\) be the number of Cs, each C has \(2\) choices (A or B), resulting in \(2^K\) possible combinations in total.
However, the constraint is \(N \le 2 \times 10^6\), and \(K\) can also be up to \(2 \times 10^6\). Therefore, an exhaustive search over all combinations takes \(O(2^N)\) time and will result in Time Limit Exceeded (TLE).
Key Observation
Whether the condition “no two As are consecutive” is satisfied depends solely on what the immediately preceding character was.
Specifically, when deciding the shift for the \(i\)-th day:
- If the previous day (day \(i-1\)) was A: day \(i\) cannot be A, so only B can be chosen.
- If the previous day (day \(i-1\)) was B: day \(i\) can freely be either A or B.
We do not need to keep track of the entire history; we only need the information of “whether the previous shift was A or B”. Thus, we can compute the answer sequentially from the beginning using Dynamic Programming (DP).
Algorithm
We iterate through day \(1\) to day \(N\) in order, maintaining only the state of the immediately preceding day.
DP State
We define the number of ways to determine the shifts up to day \(i\) as follows:
- \(dp_A\): the number of ways where the shift on day \(i\) is A
- \(dp_B\): the number of ways where the shift on day \(i\) is B
Initial State
When deciding the shift for day \(1\), it is possible to choose either A or B.
By regarding the state as “the day before was B”, we can naturally represent the state where both A and B can be chosen on day \(1\). Therefore, we initialize as follows:
- \(dp_A = 0\)
- \(dp_B = 1\)
State Transitions
Depending on the character \(c\) on day \(i\), we update the DP as follows:
When \(c =\)
A(day \(i\) must beA)- Day \(i-1\) must be
B, so the new \(dp_A\) is equal to the previous \(dp_B\). - Day \(i\) cannot be
B, so the new \(dp_B\) becomes \(0\). - \(dp_A \leftarrow dp_B\)
- \(dp_B \leftarrow 0\)
- Day \(i-1\) must be
When \(c =\)
B(day \(i\) must beB)- Day \(i-1\) can be either
AorB. - Day \(i\) cannot be
A, so the new \(dp_A\) becomes \(0\). - \(dp_A \leftarrow 0\)
- \(dp_B \leftarrow (dp_A + dp_B) \pmod{998244353}\)
- Day \(i-1\) can be either
When \(c =\)
C(day \(i\) can be eitherAorB)- When choosing
A: transitions are only possible from the state where day \(i-1\) isB. - When choosing
B: transitions are possible regardless of whether day \(i-1\) isAorB. - \(new\_dp_A \leftarrow dp_B\)
- \(dp_B \leftarrow (dp_A + dp_B) \pmod{998244353}\)
- \(dp_A \leftarrow new\_dp_A\)
- When choosing
After processing up to day \(N\), \((dp_A + dp_B) \pmod{998244353}\) is the final answer. Even if the original string already contains AA, the value will automatically be updated to \(0\), so no special corner-case handling is required.
Complexity
- Time Complexity: \(O(N)\)
- We loop through the string from beginning to end once, and each step requires only a constant number of additions and assignments. Since \(N \le 2 \times 10^6\), this easily finishes well within the time limit.
- Space Complexity: \(O(1)\) (excluding the storage for the input string)
- Since we only need to maintain the values from the previous day (day \(i-1\)), there is no need to allocate a DP array of length \(N\), and the algorithm can be implemented using only a constant number of variables.
Implementation Notes
Order of value updates: When processing \(c =\)
C, the update of \(dp_B\) uses the old value of \(dp_A\), and the update of \(dp_A\) uses the old value of \(dp_B\). To avoid overwriting values before they are used, use a temporary variable or Python’s simultaneous assignment (tuple assignment):dp_a, dp_b = dp_b, (dp_a + dp_b) % MODFast I/O: Since \(N\) can be as large as \(2 \times 10^6\), it is recommended to use
sys.stdin.readlinefor fast input reading in Python.Source Code
import sys
def main():
input = sys.stdin.readline
N = int(input())
S = input().strip()
MOD = 998244353
dp_a = 0
dp_b = 1
for c in S:
if c == "A":
dp_a = dp_b
dp_b = 0
elif c == "B":
dp_b = (dp_a + dp_b) % MOD
dp_a = 0
else:
new_a = dp_b
dp_b = (dp_a + dp_b) % MOD
dp_a = new_a
print((dp_a + dp_b) % MOD)
if __name__ == "__main__":
main()
This editorial was generated by gemini-3.8-flash-high.
posted:
last update: