Official

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

Gemini 3.8 Flash (High)

概要

A, B, C からなる長さ \(N\) の文字列の C を A または B に置き換え、A が \(2\) つ連続しないようにする組み合わせの総数を \(998244353\) で割った余りを求める問題です。

考察

各 C について A と B の \(2\) 通りの選択肢があるため、すべての置き換え方を試す素朴な全探索では最大で \(2^N\) 通りの探索が必要になります。\(N \le 2 \times 10^6\) という制約下では、この方法では実行時間制限に到底間に合わず、TLE(実行時間制限超過)となります。

そこで、「直前の日のシフトが何であったか」 だけに着目します。 \(i\) 日目のシフトを決める際、条件を満たすかどうかは直前(\(i-1\) 日目)のシフトにのみ依存します。 - \(i\) 日目を A にする場合:直前の日は AA を防ぐために必ず B でなければなりません。 - \(i\) 日目を B にする場合:直前の日は A でも B でも構いません。

このように、過去の詳細な履歴すべてを保持する必要はなく、「直前が A であった場合の数」と「直前が B であった場合の数」だけを管理すればよいことが分かります。これは動的計画法(DP)を用いることで効率よく計算できます。

アルゴリズム

動的計画法を用いて、\(1\) 日目から順にシフトの割り当て数を計算します。

DPテーブルの定義

\(i\) 日目まで決めた時点での状態を以下のように定義します。 - \(\text{dp\_a}\) : \(i\) 日目のシフトが A であるような割り当て方の総数 - \(\text{dp\_b}\) : \(i\) 日目のシフトが B であるような割り当て方の総数

初期状態(1日目)

\(S\) の \(1\) 文字目 \(S[0]\) に応じて初期化します。 - \(S[0] =\) 'A' の場合: \(\text{dp\_a} = 1, \text{dp\_b} = 0\) - \(S[0] =\) 'B' の場合: \(\text{dp\_a} = 0, \text{dp\_b} = 1\) - \(S[0] =\) 'C' の場合: \(\text{dp\_a} = 1, \text{dp\_b} = 1\)

遷移(\(i = 2, 3, \dots, N\) 日目)

次の日の状態を \(\text{next\_dp\_a}, \text{next\_dp\_b}\) とします(初期値は \(0\))。

  1. \(i\) 日目を A にできる場合(\(S[i-1]\) が 'A' または 'C'):
    • 直前は B である必要があるため、\(\text{next\_dp\_a} = \text{dp\_b}\)
  2. \(i\) 日目を B にできる場合(\(S[i-1]\) が 'B' または 'C'):
    • 直前は A でも B でもよいため、\(\text{next\_dp\_b} = \text{dp\_a} + \text{dp\_b}\)

遷移後、\(\text{dp\_a} \leftarrow \text{next\_dp\_a}\)、\(\text{dp\_b} \leftarrow \text{next\_dp\_b}\) と更新します。

答え

すべての日の遷移を終えた後の \(\text{dp\_a} + \text{dp\_b}\) が答えとなります。なお、計算はすべて \(998244353\) を法として行います。

計算量

  • 時間計算量: \(O(N)\)
    • 各日について定数時間(\(O(1)\))で遷移を計算できるため、全体で \(O(N)\) となり実行時間制限に十分間に合います。
  • 空間計算量: \(O(N)\)
    • 入力文字列を保持するために \(O(N)\) のメモリを使用します。DP の状態変数は直前の値のみを保持するため \(O(1)\) です。

実装のポイント

  • 状態の再利用: 配列としてサイズ \(N\) の DP テーブルを確保しても間に合いますが、直前の日の状態(変数 \(2\) つ)だけを持っておくことで、メモリ消費量を抑えてシンプルに実装できます。

  • ModInt の利用: 足し算のたびに余りを計算する必要があります。AtCoder Library の atcoder::modint998244353 などを利用すると、オーバーフローや余りの取り忘れを防ぎ、コードを簡潔に書くことができます。

    ソースコード

#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;
}

この解説は gemini-3.8-flash-high によって生成されました。

posted:
last update: