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: