Official

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:

  1. When \(c =\) A (day \(i\) must be A)

    • 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\)
  2. When \(c =\) B (day \(i\) must be B)

    • Day \(i-1\) can be either A or B.
    • 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}\)
  3. When \(c =\) C (day \(i\) can be either A or B)

    • When choosing A: transitions are only possible from the state where day \(i-1\) is B.
    • When choosing B: transitions are possible regardless of whether day \(i-1\) is A or B.
    • \(new\_dp_A \leftarrow dp_B\)
    • \(dp_B \leftarrow (dp_A + dp_B) \pmod{998244353}\)
    • \(dp_A \leftarrow new\_dp_A\)

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) % MOD
    

  • Fast I/O: Since \(N\) can be as large as \(2 \times 10^6\), it is recommended to use sys.stdin.readline for 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: