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 = 0dpB = 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: