ABC215-E Chain Contestant 解説
文字列に含まれる文字の種類数を $σ$ とおきます。 $O(N2^{σ})$ で解きます。
問題
リンクはこちら
解法
$σ = 1$
まずは $σ = 1$ の場合を考えてみましょう。このとき答えは明らかに $2^{N} - 1$ です。これをDPで解いてみることにします。配列 $g[i] \coloneqq \text{$s[0, i)$ の範囲についての選び方}$ を考えます。例えば $s = \text{\textquotedblleft AAAAAA\textquotedblright}$ の時、 $g = [0, 1, 3, 7, 15, 31, 63]$ です。
ここで、$g[i]$ の値を求める計算を、以下のように場合分けして考えてみましょう。
最初に選ぶのが $s[0]$ のとき
$s[1], s[2], \ldots, s[i-1]$ の $i-1$ 個については、選んでも選ばなくてもよい。
$\Rightarrow 2^{i-1}$ 通り
最初に選ぶのが $s[1]$ のとき
$s[2], s[3], \ldots, s[i-1]$ の $i-2$ 個については、選んでも選ばなくてもよい。
$\Rightarrow 2^{i-2}$ 通り
$\vdots$
最初に選ぶのが $s[i-1]$ のとき
$1$ 通り
これらを足し合わせることで、確かに $\displaystyle\sum _ {j=0}^{i-1}2^{j} = 2^{i} - 1$ となり、 $g[i]$ の値と一致します。
$σ = 2$
続いて、 $σ = 2$ の場合を考えます。問題の文字列 $s$ が $\text{\textquotedblleft ABAABAB\textquotedblright}$ であるとします。 同じ種類のコンテストは必ず一かたまりで選ばなければなりませんが、この順番が $\text{A} \rightarrow \text{B}$ の場合を考えます。
まず、先ほど考えた $σ = 1$ のパターン同様、 $\text{A}$ だけを選ぶ場合を表す配列 $g _ {A}$ が存在すると考えます。これは以下のようになります。
では、改めて $g _ {\text{AB}}[i] \coloneqq \text{$s[0, i)$ から条件を満たすように A, B の順で選ぶ方法の個数}$ を埋めていきます。 $g _ {\text{AB}}[7]$ を計算するときの様子を考えてみます。
最初に選ぶ $\text{B}$ が $s[1]$ のとき
$s[0, 1)$ からいくつか $\text{A}$ を選び、$s[1]$ を選ぶ。$s[4], s[6]$ は選んでも選ばなくてもよい。
$\Rightarrow g _ {\text{A}}[2] \times 2^{2}$ 通り
最初に選ぶ $\text{B}$ が $s[4]$ のとき
$s[0, 4)$ からいくつか $\text{A}$ を選び、$s[4]$ を選ぶ。$s[6]$ は選んでも選ばなくてもよい。
$\Rightarrow g _ {\text{A}}[4] \times 2^{1}$ 通り
最初に選ぶ $\text{B}$ が $s[6]$ のとき
$s[0, 6)$ からいくつか $\text{A}$ を選び、$s[6]$ を選ぶ。
$\Rightarrow g _ {\text{A}}[6] \times 2^{0}$ 通り
これらの合計が $g _ {\text{AB}}[7]$ になります。足し合わせるべき $g _ {\text{A}}[i]$ を以下のように累積的に計算すれば、前から見ていくことで $g _ {\text{AB}}$ を計算することができます。
acc = 0 for i in range(N): if s[i] == 'B': acc = acc * 2 + gA[i] gAB[i+1] = acc else: gAB[i+1] = gAB[i]
同様に $g _ {\text{B}}$ から $g _ {\text{BA}}$ も計算できます。
ここで、 $g _ {\text{\{A,B\}}}[i] \coloneqq \text{$s[0, i)$ から条件を満たすように A, B を選ぶ方法の個数}$ という配列を考えると、 $g _ {\text{\{A,B\}}}[i] = g _ {\text{AB}}[i] + g _ {\text{BA}}[i]$ が成り立ちます。また、文字列中に出現する文字全体の集合を $U$ とおくと、本問題で求めるべき答えは
$$ \sum _ {\substack{S \subset U \\ S \neq \varnothing}} g _ {S}[n] $$
と表せます。
$σ \ge 3$ のとき
これまでとおおむね同様にして $g _ {*}$ を求めていくことができます。 例えば $g _ {\text{ABCD}}$ であれば以下のようなコードになります。
acc = 0 for i in range(N): if s[i] == 'D': acc = acc * 2 + gABC[i] gABCD[i+1] = acc else: gABCD[i+1] = gABCD[i]
特に $σ$ が大きくなると、ありうるすべての順列について $g _ {*}$ を足し合わせるのが困難となるので、直接 $g _ {S}$ を求める必要があります。この方法について、 $σ = 2$ の場合に戻って考えます。
$g _ {S}$ を直接求める
$s = \text{\textquotedblleft ABAABAB\textquotedblright}$ とします。ここで、以下の二種類の配列を考えます。
$$ \begin{aligned} f _ {\text{AB}}[i] &\coloneqq \text{$s[i-1]$ を必ず選ぶとき、 $s[0, i)$ から条件を満たすように A, B の順で選ぶ方法の個数} \\ f _ {\text{\{A,B\}}}[i] &\coloneqq \text{$s[i-1]$ を必ず選ぶとき、 $s[0, i)$ から条件を満たすように A, B を選ぶ方法の個数} \end{aligned} $$
この時、 $f _ {\text{\{A,B\}}}[i] = f _ {\text{AB}}[i] + f _ {\text{BA}}[i]$ です。また、条件を満たす選び方それぞれについて、最後の文字の位置がどこであるかを考えると、$g$ が $f$ の累積和となっていることがわかります。
さらに、 $f _ {\text{AB}}$ と $f _ {\text{BA}}$ を見比べると、一方が $0$ でないときもう一方は必ず $0$ となっています。これは、 $s[i-1]$ を必ず選ぶという条件により、$s[i-1] = \text{A}$ のときは $f _ {\text{AB}}[i] = 0$ となるためです。
$f _ {\text{AB}}$ は以下のように計算可能ですが、
fAB = [0] * (N+1) acc = 0 for i in range(N): if s[i] == 'B': # s[i-1] を必ず使うという条件により、2の指数が1減る # それに伴い更新式が若干変化した fAB[i+1] = acc + g_A_[i] acc = acc * 2 + g_A_[i]
$f _ {\text{AB}}$ と $f _ {\text{BA}}$ の更新のタイミングはかぶらないため、if, else によって以下のようにまとめることができます。
fAB = [0] * (N+1) fBA = [0] * (N+1) accA = 0 accB = 0 for i in range(N): if s[i] == 'A': fBA[i+1] = accA + g_B_[i] accA = accA * 2 + g_B_[i] else: fAB[i+1] = accB + g_A_[i] accB = accB * 2 + g_A_[i]
さらに配列もひとつにまとめてしまうことで、直接 $f _ {\text{\{A,B\}}}$ を計算することができます。
f_AB_= [0] * (N+1) accA = 0 accB = 0 for i in range(N): if s[i] == 'A': f_AB_[i+1] = accA + g_B_[i] accA = accA * 2 + g_B_[i] else: f_AB_[i+1] = accB + g_A_[i] accB = accB * 2 + g_A_[i]
以上の内容と、 $f$ の累積和を取ると $g$ になることから、$g _ {\text{\{A,B\}}}$ は結局以下のように直接計算することができます。
g_AB_= [0] * (N+1) accA = 0 accB = 0 for i in range(N): if s[i] == 'A': g_AB_[i+1] = accA + g_B_[i] accA = accA * 2 + g_B_[i] else: g_AB_[i+1] = accB + g_A_[i] accB = accB * 2 + g_A_[i] g_AB_[i+1] += g_AB_[i]
一般の $σ$ について
より多くの文字が含まれる場合についても、最後に選ぶ文字がどれであるかによって分類すると、同様の計算が可能です。 $\mathrm{dp}[S][i] \coloneqq \text{$s[0, i)$ から条件を満たすように、 $S$ 中の文字を選ぶ方法の個数}$ を埋めていきます。 更新式は以下のように表されます。
# 更新のイメージ for i in range(N): c = s[i] if c in S: dp[S][i+1] = acc[c] + dp[S\{c}][i] acc[c] = acc[c] * 2 + dp[S\{c}][i] dp[S][i+1] += dp[S][i]
ここで、 $\mathrm{dp}[S]$ を計算する際には $\mathrm{dp}[S \setminus \{c\}]$ が既に計算されている必要があることに注意してください。
実際に実装する際には $S$ はビット表現として持つことになるので、以下のようなコードになります。
# 各文字は適当な整数に変換 for i in range(N): c = s[i] if bit>>c & 1: dp[bit][i+1] = acc[c] + dp[bit ^ 1<<c][i] acc[c] = acc[c] * 2 + dp[bit ^ 1<<c][i] dp[bit][i+1] += dp[bit][i]
bit について昇順に更新していけば、$\mathrm{dp}[S]$ の更新に必要な $\mathrm{dp}[T] ~ (T \subset S)$ は全て計算済みとなります。
余談
逆に言えば自分の上位集合はまだ計算されていないことになるので、if を省略することができます。
for i in range(N): # bit & 1<<s[i] == 0 なら dp[bit | 1<<s[i]][i] == 0 # acc[s[i]] もずっと 0 のまま dp[bit][i+1] = acc[s[i]] + dp[bit ^ 1<<s[i]][i] + dp[bit][i] acc[s[i]] = acc[s[i]] * 2 + dp[bit ^ 1<<s[i]][i]
実装
全体の実装は以下の通りです。計算量は $O(N2^σ)$ です。
K = 10 MOD = 998_244_353 N = int(input()) s = [ord(x) - ord('A') for x in input()] dp = [[0] * (N+1) for _ in range(1<<K)] dp[0][0] = 1 for bit in range(1<<K): acc = [0] * K for i in range(N): dp[bit][i+1] = (acc[s[i]] + dp[bit ^ 1<<s[i]][i] + dp[bit][i]) % MOD acc[s[i]] = (acc[s[i]] * 2 + dp[bit ^ 1<<s[i]][i]) % MOD ans = sum(dp[bit][N] for bit in range(1, 1<<K)) % MOD print(ans)
