Official

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

gpt-6-astra-medium

概要

未定のシフト C を A または B に置き換え、AA が現れない割り当て方の数を求めます。左から順に文字を処理し、末尾が A か B かによって場合の数を管理する動的計画法を使います。

考察

C が \(K\) 個あるとき、置き換え方は全部で \(2^K\) 通りあります。すべてを試す方法では、\(N\) が最大 \(2 \times 10^6\) なので間に合いません。

ここで、条件を満たす文字列の末尾に、新しく文字を追加することを考えます。

  • A を追加できるのは、直前の文字が B の場合だけです。
  • B は、直前の文字が A でも B でも追加できます。

つまり、次の文字を追加できるかどうかの判断には、それまでの文字列全体ではなく、末尾の文字だけが必要です。

そこで、処理済みの部分について、条件を満たす割り当て方を次の \(2\) 種類に分けて数えます。

  • \(a\):末尾が A である割り当て方の数
  • \(b\):末尾が B である割り当て方の数

この \(2\) つを更新しながら進めれば、すべての置き換えを列挙する必要はありません。

アルゴリズム

1. 初期化

まだ文字を処理していない状態を、

\(a = 0,\quad b = 1\)

とします。

空の文字列に末尾の文字はありませんが、ここでは「先頭に A も B も置ける \(1\) 通りの状態」を、便宜上 \(b\) に含めています。文字列の直前に仮想的な B があると考えても構いません。

2. 左から順に更新する

更新後の値を \(a', b'\) とすると、現在の文字に応じた遷移は次のようになります。

現在の文字 \(a'\) \(b'\)
A \(b\) \(0\)
B \(0\) \(a+b\)
C \(b\) \(a+b\)

それぞれの理由は以下の通りです。

現在の文字が A

A を置くことが確定しています。AA を避けるため、直前が B の割り当てだけを引き継ぎます。

\(a' = b,\quad b' = 0\)

現在の文字が B

B を置くことが確定しています。直前がどちらの文字でもよいため、両方の場合の数を足します。

\(a' = 0,\quad b' = a+b\)

現在の文字が C

A と B の両方を選べます。

  • A にする場合は、直前が B の割り当てだけを引き継ぎます。
  • B にする場合は、すべての割り当てを引き継ぎます。

\(a' = b,\quad b' = a+b\)

加算結果は、その都度 \(998244353\) で割った余りにします。

3. 答えを求める

すべての文字を処理した後、末尾が A の場合と B の場合を合計します。

答えは \((a+b) \bmod 998244353\) です。

各遷移では、AA を作らない追加だけを行っています。また、条件を満たす割り当ては必ずいずれかの遷移で数えられるため、漏れや重複なく総数を求められます。

具体例

\(S = \texttt{CCA}\) の場合を考えます。

処理済みの部分 \(a\) \(b\)
空文字列 \(0\) \(1\)
C \(1\) \(1\)
CC \(1\) \(2\)
CCA \(2\) \(0\)

条件を満たす文字列は ABA と BBA の \(2\) 通りなので、答えは \(2\) です。

計算量

  • 時間計算量: \(O(N)\)
  • 空間計算量: \(O(N)\)。入力文字列の保存に必要です。DP のための追加領域は、変数を数個使うだけなので \(O(1)\) です。

実装のポイント

  • 更新には更新前の値を使います。
    例えば C の処理では、a, b = b, total と同時代入することで、更新前の \(b\) を新しい \(a\) に設定できます。

  • 剰余は引き算 \(1\) 回で求められます。
    常に \(0 \le a,b < \mathrm{MOD}\) を保っているため、\(a+b < 2\mathrm{MOD}\) です。したがって、和が \(\mathrm{MOD}\) 以上なら \(\mathrm{MOD}\) を \(1\) 回引くだけで十分です。

  • 入力文字列を bytes として処理しています。
    sys.stdin.buffer.readline で読み込んだ文字列を走査すると、各要素は整数になります。そのため、コードでは A を 65、B を 66 と比較しています。

  • 不可能な場合の特別処理は不要です。
    例えば確定した AA があると、遷移によって \(a=b=0\) になります。それ以降も両方が \(0\) のままなので、自然に答えとして \(0\) が出力されます。

    ソースコード

import sys

def main():
    input = sys.stdin.buffer.readline
    n = int(input())
    s = input().strip()

    MOD = 998244353
    a, b = 0, 1

    for c in s:
        if c == 65:  # A
            a, b = b, 0
        else:
            total = a + b
            if total >= MOD:
                total -= MOD
            if c == 66:  # B
                a, b = 0, total
            else:  # C
                a, b = b, total

    answer = a + b
    if answer >= MOD:
        answer -= MOD
    print(answer)

if __name__ == "__main__":
    main()

この解説は gpt-6-astra-medium によって生成されました。

posted:
last update: