Official

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\) に応じて、次のように更新します。

  1. \(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\)
  2. \(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}\)
  3. \(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: