D - シフトの割り当て / Shift Assignment Editorial by admin
Gemini 3.8 Flash (High)概要
長さ \(N\) の文字列 \(S\) に含まれる C を A または B に置き換えることで、「A が \(2\) つ連続しない」という条件を満たす文字列が何通り作れるかを \(998244353\) で割った余りで求める問題です。
考察
素朴なアプローチ
C の個数を \(K\) とすると、各 C について A か B の \(2\) 通りの選択肢があるため、全体では \(2^K\) 通りの組み合わせが考えられます。
しかし、制約は \(N \le 2 \times 10^6\) であり、\(K\) も最大で \(2 \times 10^6\) となるため、すべてを全探索すると \(O(2^N)\) となり実行時間制限(TLE)に間に合いません。
重要な気づき
「A が連続しない」という制約を満たすかどうかは、直前の文字が何であったか のみに依存します。
具体的には、\(i\) 日目のシフトを決めるとき:
- 直前(\(i-1\) 日目)が A の場合:\(i\) 日目は A にできず、B のみ が選べます。
- 直前(\(i-1\) 日目)が B の場合:\(i\) 日目は A でも B でも 自由に選べます。
過去の全履歴を覚えておく必要はなく、「直前のシフトが A か B か」という情報だけを持てばよいため、動的計画法(DP) を用いて前から順に計算していくことができます。
アルゴリズム
直前の状態のみを保持しながら、1日目から \(N\) 日目まで順番に見ていきます。
DP の状態
\(i\) 日目まで決定した時点での場合の数を以下のように定義します。
- \(dp_A\): \(i\) 日目のシフトが A である場合の数
- \(dp_B\): \(i\) 日目のシフトが B である場合の数
初期状態
1日目のシフトを決める際、1日目を A にすることも B にすることも可能です。
「直前の日が B であった」と見なすことで、1日目に A も B も選べる状態を自然に表現できるため、以下のように初期化します。
- \(dp_A = 0\)
- \(dp_B = 1\)
状態遷移
\(i\) 日目の文字 \(c\) に応じて、次のように更新します。
\(c =\)
Aの場合(\(i\) 日目をAにしなければならない)- \(i-1\) 日目は
Bでなければならないため、新しい \(dp_A\) は直前の \(dp_B\) と等しくなります。 - \(i\) 日目を
Bにすることはできないため、新しい \(dp_B\) は \(0\) になります。 - \(dp_A \leftarrow dp_B\)
- \(dp_B \leftarrow 0\)
- \(i-1\) 日目は
\(c =\)
Bの場合(\(i\) 日目をBにしなければならない)- \(i-1\) 日目は
AでもBでも構いません。 - \(i\) 日目を
Aにすることはできないため、新しい \(dp_A\) は \(0\) になります。 - \(dp_A \leftarrow 0\)
- \(dp_B \leftarrow (dp_A + dp_B) \pmod{998244353}\)
- \(i-1\) 日目は
\(c =\)
Cの場合(\(i\) 日目はAでもBでもよい)Aを選ぶ場合:\(i-1\) 日目がBの状態からのみ遷移できます。Bを選ぶ場合:\(i-1\) 日目がAでもBでも遷移できます。- \(new\_dp_A \leftarrow dp_B\)
- \(dp_B \leftarrow (dp_A + dp_B) \pmod{998244353}\)
- \(dp_A \leftarrow new\_dp_A\)
\(N\) 日目まで走査した後の \((dp_A + dp_B) \pmod{998244353}\) が答えとなります。元の文字列に AA が含まれている場合なども、自動的に値が \(0\) に更新されるため例外処理は不要です。
計算量
- 時間計算量: \(O(N)\)
- 文字列の先頭から末尾まで \(1\) 回ループを回し、各ステップで行う処理は定数回の加算・代入のみです。\(N \le 2 \times 10^6\) のため、実行時間制限内に十分余裕を持って間に合います。
- 空間計算量: \(O(1)\)(入力文字列の保持を除いた場合)
- 前の日(\(i-1\) 日目)の値のみを保持すればよいため、DP配列を長さ \(N\) で確保する必要がなく、定数個の変数のみで実装できます。
実装のポイント
- 値の更新順序:
\(c =\)
Cの処理では、\(dp_B\) の更新に古い \(dp_A\) の値を用い、\(dp_A\) の更新には古い \(dp_B\) の値を用います。値を上書きしてしまわないよう、一時変数を使うか、Pythonの同時代入(タプル代入)を利用して更新します。 “`python dp_a, dp_b = dp_b, (dp_a + dp_b) % MOD
- **高速な入出力**:
$N$ が最大で $2 \times 10^6$ と大きいため、Pythonでは `sys.stdin.readline` を用いて高速に入力を受け取ることが望ましいです。
## ソースコード
```python
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()
この解説は gemini-3.8-flash-high によって生成されました。
posted:
last update: