Official

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

gpt-6-astra-medium

概要

C を A または B に置き換え、AA を含まない文字列を作る方法の数を求めます。左から順に文字を処理し、末尾が A か B かで場合分けする動的計画法(DP)を使います。

考察

すべての置き換えを試すと間に合わない

C が \(K\) 個ある場合、置き換え方は \(2^K\) 通りあります。これらをすべて列挙して AA がないか調べる方法では、\(N\) が最大 \(2 \times 10^6\) なので間に合いません。

必要なのは直前の文字だけ

すでに条件を満たしている文字列の末尾に、次の文字を追加することを考えます。

  • A を追加する場合、直前の文字は B でなければなりません。
  • B を追加する場合、直前の文字は A でも B でも構いません。

つまり、次の文字を追加できるかの判定には、それまでの文字列全体ではなく、末尾の文字だけが必要です。

そこで、条件を満たす置き換え方を「末尾が A のもの」と「末尾が B のもの」に分けて数えます。

アルゴリズム

DP の状態

先頭から何文字かを処理した時点で、次の \(2\) つの値を管理します。値はすべて \(998244353\) で割った余りとします。

  • dpA:処理済みの部分が条件を満たし、末尾が A となる置き換え方の数
  • dpB:処理済みの部分が条件を満たし、末尾が B となる置き換え方の数

初期化

まだ文字を処理していない状態を、

  • dpA = 0
  • dpB = 1

とします。

空文字列の置き換え方は \(1\) 通りです。先頭には A も B も置けるため、文字列の直前に仮想的な B があると考えて dpB にこの \(1\) 通りを持たせます。

遷移

現在の文字を c とし、処理後の値を nextA, nextB とします。

現在の位置に A を置く場合

c が A または C なら A を置けます。ただし、直前が A だと AA ができるので、末尾が B の場合からしか遷移できません。

  • c != 'B' なら、\(\mathrm{nextA} = \mathrm{dpB}\)
  • それ以外なら、\(\mathrm{nextA} = 0\)

現在の位置に B を置く場合

c が B または C なら B を置けます。直前の文字に制限はありません。

  • c != 'A' なら、\(\mathrm{nextB} = (\mathrm{dpA} + \mathrm{dpB}) \bmod 998244353\)
  • それ以外なら、\(\mathrm{nextB} = 0\)

両方を計算した後、dpA, dpB をそれぞれ nextA, nextB に更新します。

この遷移は、新たに AA を作らない追加方法をすべて数えています。また、追加する文字が A の場合と B の場合は重複しないため、各置き換え方をちょうど \(1\) 回ずつ数えられます。

具体例

\(S = \texttt{CCA}\) の場合を考えます。

処理済みの部分 dpA dpB
空文字列 \(0\) \(1\)
C \(1\) \(1\)
CC \(1\) \(2\)
CCA \(2\) \(0\)

CC まででは、末尾が A のものは BA、末尾が B のものは AB, BB です。

最後の A を追加できるのは末尾が B の場合だけなので、完成する文字列は ABA, BBA の \(2\) 通りです。

答え

すべての文字を処理した後、末尾が A の場合と B の場合を合計します。

答えは \((\mathrm{dpA} + \mathrm{dpB}) \bmod 998244353\) です。条件を満たす置き換え方がなければ、両方が \(0\) となり、答えも \(0\) になります。

計算量

  • 時間計算量: \(O(N)\)
    各文字について定数回の計算を行います。
  • 空間計算量: \(O(N)\)
    提示コードでは入力文字列を保存するために \(O(N)\) の領域を使います。DP 自体の追加領域は \(O(1)\) です。

実装のポイント

  • 更新前の値を使って両方の遷移を計算します。
    dpA を先に書き換えると、nextB の計算に更新後の値を使ってしまいます。提示コードのように nextA, nextB を用意してから更新しましょう。

  • 各状態を常に \(998244353\) 未満に保ちます。
    dpA + dpB は最大でも \(2 \times (998244353 - 1)\) なので、法以上なら \(1\) 回引くだけで余りを求められます。この和は C++ の符号付き \(32\) ビット整数の範囲にも収まります。

  • 固定された AA を別途検出する必要はありません。
    A の直後に A を処理すると、dpA, dpB がどちらも \(0\) になるため、DP だけで不可能な場合も正しく扱えます。

    ソースコード

#include <iostream>
#include <string>
using namespace std;

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int N;
    string S;
    cin >> N >> S;

    constexpr int MOD = 998244353;
    int dpA = 0, dpB = 1;

    for (char c : S) {
        int nextA = (c != 'B') ? dpB : 0;
        int nextB = 0;
        if (c != 'A') {
            nextB = dpA + dpB;
            if (nextB >= MOD) nextB -= MOD;
        }
        dpA = nextA;
        dpB = nextB;
    }

    int answer = dpA + dpB;
    if (answer >= MOD) answer -= MOD;
    cout << answer << '\n';

    return 0;
}

この解説は gpt-6-astra-medium によって生成されました。

posted:
last update: