Official

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

Gemini 3.8 Flash (High)

Overview

This problem asks us to find the number of ways to replace each C with either A or B in a string of length \(N\) consisting of A, B, and C, such that no two As are consecutive, modulo \(998244353\).

Analysis

Since there are \(2\) choices (A or B) for each C, a naive brute-force search trying all replacements would require exploring up to \(2^N\) possibilities. Under the constraint \(N \le 2 \times 10^6\), this approach will not finish within the time limit and will result in TLE (Time Limit Exceeded).

Therefore, we focus only on “what the previous day’s shift was.” When deciding the shift for the \(i\)-th day, whether the condition is satisfied depends only on the shift of the previous day (\((i-1)\)-th day). - If the \(i\)-th day is set to A: To prevent AA, the previous day must be B. - If the \(i\)-th day is set to B: The previous day can be either A or B.

Thus, we do not need to keep track of the entire past history; we only need to maintain “the number of ways where the previous day was A” and “the number of ways where the previous day was B.” This can be computed efficiently using Dynamic Programming (DP).

Algorithm

Using dynamic programming, we compute the number of valid shift assignments sequentially from the \(1\)st day.

DP State Definition

We define the state after deciding up to the \(i\)-th day as follows: - \(\text{dp\_a}\): Total number of valid assignments such that the \(i\)-th day’s shift is A - \(\text{dp\_b}\): Total number of valid assignments such that the \(i\)-th day’s shift is B

Initial State (Day 1)

Initialize according to the first character of \(S\), \(S[0]\): - If \(S[0] =\) 'A': \(\text{dp\_a} = 1, \text{dp\_b} = 0\) - If \(S[0] =\) 'B': \(\text{dp\_a} = 0, \text{dp\_b} = 1\) - If \(S[0] =\) 'C': \(\text{dp\_a} = 1, \text{dp\_b} = 1\)

Transitions (Days \(i = 2, 3, \dots, N\))

Let the state for the next day be \(\text{next\_dp\_a}\) and \(\text{next\_dp\_b}\) (initialized to \(0\)).

  1. If the \(i\)-th day can be A (\(S[i-1]\) is 'A' or 'C'):
    • Since the previous day must be B, \(\text{next\_dp\_a} = \text{dp\_b}\).
  2. If the \(i\)-th day can be B (\(S[i-1]\) is 'B' or 'C'):
    • Since the previous day can be either A or B, \(\text{next\_dp\_b} = \text{dp\_a} + \text{dp\_b}\).

After computing the transitions, update \(\text{dp\_a} \leftarrow \text{next\_dp\_a}\) and \(\text{dp\_b} \leftarrow \text{next\_dp\_b}\).

Answer

After completing the transitions for all days, the answer is \(\text{dp\_a} + \text{dp\_b}\). All calculations should be performed modulo \(998244353\).

Complexity

  • Time Complexity: \(O(N)\)
    • Since the transition for each day can be computed in constant time (\(O(1)\)), the overall time complexity is \(O(N)\), which easily runs within the time limit.
  • Space Complexity: \(O(N)\)
    • \(O(N)\) memory is used to store the input string. The DP state variables only need to retain the values from the previous step, requiring \(O(1)\) auxiliary space.

Implementation Notes

  • Reusing State: While allocating an array of size \(N\) for the DP table would also pass within limits, maintaining only the previous day’s state (just \(2\) variables) reduces memory usage and simplifies the implementation.

  • Using ModInt: A modulo operation is required after each addition. Using utilities such as atcoder::modint998244353 from the AtCoder Library helps prevent integer overflow and missed modulo operations, keeping the code clean and concise.

    Source Code

#include <iostream>
#include <string>
#include <atcoder/modint>

using namespace std;
using mint = atcoder::modint998244353;

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    int n;
    if (!(cin >> n)) return 0;
    string s;
    cin >> s;

    mint dp_a = 0;
    mint dp_b = 0;

    if (s[0] == 'A') {
        dp_a = 1;
    } else if (s[0] == 'B') {
        dp_b = 1;
    } else { // 'C'
        dp_a = 1;
        dp_b = 1;
    }

    for (int i = 1; i < n; ++i) {
        mint next_dp_a = 0;
        mint next_dp_b = 0;

        if (s[i] == 'A' || s[i] == 'C') {
            next_dp_a = dp_b;
        }
        if (s[i] == 'B' || s[i] == 'C') {
            next_dp_b = dp_a + dp_b;
        }

        dp_a = next_dp_a;
        dp_b = next_dp_b;
    }

    mint ans = dp_a + dp_b;
    cout << ans.val() << "\n";

    return 0;
}

This editorial was generated by gemini-3.8-flash-high.

posted:
last update: