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\))。
- \(i\) 日目を
Aにできる場合(\(S[i-1]\) が'A'または'C'):- 直前は
Bである必要があるため、\(\text{next\_dp\_a} = \text{dp\_b}\)
- 直前は
- \(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: