Official

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

Claude Opus 5 (High)

概要

A, B, C からなる文字列 \(S\) の C をそれぞれ A または B に置き換えて、AA を含まない文字列を作る方法の総数を \(998244353\) で割った余りを求める問題です。典型的な「連続禁止」条件の数え上げで、桁 DP(1 文字ずつ進める DP)で解けます。

考察

素朴なアプローチ

C の個数を \(k\) とすると、置き換え方は \(2^k\) 通りあります。\(N \le 2\times 10^6\) なので \(k\) も最大 \(2\times 10^6\) となり、全列挙は到底不可能です(\(2^{2000000}\) 通り)。

重要な気づき

AA を含まないという条件は「隣り合う 2 文字がともに A にならない」という局所的な条件です。つまり、\(i\) 文字目を決めるとき、条件違反が起きるかどうかは「\(i-1\) 文字目が A だったか否か」だけで判定できます。

したがって、先頭から順に文字を確定していくとき、覚えておくべき情報は

  • ここまでの文字列が条件を満たしている
  • 直前の文字が A かどうか

の 2 状態だけです。「ここまで何通りあるか」を状態ごとに持って足し合わせていけば、指数的な列挙を線形時間の DP に落とせます。

アルゴリズム

以下の DP を考えます。

  • \(a_i\) : 先頭 \(i\) 文字を条件を満たすように決め、\(i\) 文字目が A である場合の通り数
  • \(b_i\) : 先頭 \(i\) 文字を条件を満たすように決め、\(i\) 文字目が B である場合の通り数

初期値は「\(0\) 文字目は A でない」と考えて \(a_0 = 0,\ b_0 = 1\) とします(\(b\) を「直前が A でない状態」と解釈すると自然です)。

\(i\) 文字目 \(S_i\) による遷移は次の通りです。

  1. \(S_i = \) A のとき:\(i\) 文字目は必ず A。直前が A だと AA になるので不可、直前が B(または先頭)のときだけ許される。 $\(a_i = b_{i-1},\qquad b_i = 0\)$

  2. \(S_i = \) B のとき:\(i\) 文字目は必ず B。直前は何でもよい。 $\(a_i = 0,\qquad b_i = a_{i-1} + b_{i-1}\)$

  3. \(S_i = \) C のとき:A にする場合と B にする場合の両方を足す(上の 2 つを合わせたもの)。 $\(a_i = b_{i-1},\qquad b_i = a_{i-1} + b_{i-1}\)$

答えは \(a_N + b_N \bmod 998244353\) です。条件を満たす方法が存在しなければ自動的に \(0\) になります(例:\(S = \) AA なら \(a_1=1, b_1=0 \Rightarrow a_2 = b_1 = 0, b_2 = 0\))。

具体例

\(S = \) CAC の場合:

\(i\) 文字 \(a_i\) \(b_i\)
0 - 0 1
1 C \(b_0=1\) \(a_0+b_0=1\)
2 A \(b_1=1\) 0
3 C \(b_2=0\) \(a_2+b_2=1\)

答えは \(0 + 1 = 1\)。実際 BAB のみが条件を満たします(AAB, AAA, BAA はすべて AA を含む)。

計算量

  • 時間計算量: \(O(N)\)
  • 空間計算量: \(O(N)\)(入力文字列の保持分。DP 自体は \(O(1)\))

実装のポイント

  • 配列を持たない:\(a_i, b_i\) は直前の値だけから決まるので、変数 2 つを使い回します。\(N = 2\times 10^6\) なので長さ \(N\) の配列を作るとメモリ・時間ともに無駄です。Python では a, b = b, (a+b) % M のようにタプル代入すると一時変数不要で書けます。

  • mod を毎回取る:\(a+b\) が巨大化しないよう毎ステップ \(998244353\) で剰余を取ります(Python では多倍長整数のため、取らないと桁数爆発で極端に遅くなります)。

  • 早期終了:\(a = b = 0\) になった時点でそれ以降どう遷移しても \(0\) のままなので、0 を出力して即終了できます。AA が確定で含まれるケースを高速に処理できます。

  • 入力の読み込み:\(N\) が大きいので sys.stdin.buffer.read() でまとめて読み込みます。バイト列を直接 for ch in s でループすると各要素は整数(A は 65、B は 66)になるため、文字比較より高速に判定できます。

  • \(N \ge 1\) でも \(S\) が空になる入力形式のズレに備えて、len(data) > 1 のチェックを入れておくと安全です。

    ソースコード

import sys

def main():
    data = sys.stdin.buffer.read().split()
    n = int(data[0])
    s = data[1] if len(data) > 1 else b""
    M = 998244353
    a = 0  # ends with A
    b = 1  # ends with B (or start)
    A = 65
    B = 66
    for ch in s:
        if ch == A:
            a, b = b, 0
        elif ch == B:
            a, b = 0, (a + b) % M
        else:
            a, b = b, (a + b) % M
        if b == 0 and a == 0:
            print(0)
            return
    print((a + b) % M)

main()

この解説は claude5opus-high によって生成されました。

posted:
last update: