ABC167E-Colorful BlocksをDPで解きたかった話
問題概要
一列に$N$個のブロックが置いてある。それぞれのブロックを色$1$~$ M $のいずれかで塗るとき，隣接かつ同色なブロックの組が$K$個以下になるような塗り方の総数を$mod998244353$で求めよ。
コンテスト中
数え上げが苦手すぎて$O(N^2 M^2)$の$dp$が浮かんだ瞬間に飛ばしてFに行きました。(ちなみにFは結局解けずに冷えました。)
動的計画法で解きたい
まず，問題文を読むと次のような$dp$が浮かびます。
$dp[x][y][z]:=$先頭$x$個のブロックを隣接かつ同色な組が$y$組になるように，$x$個目のブロックを色$z$で塗るような塗り方の総数
このとき，遷移は
\begin{eqnarray} dp[x][y][z] = \begin{cases} 1 & ( x = 1 かつ y = 0 ) \\ dp[x-1][y-1][z] - dp[x-1][y][z] + \displaystyle \sum_{w=1}^{M} dp[x-1][y][w] & ( x > 1 かつ x > y ) \\ 0 & (otherwise) \\ \end{cases} \end{eqnarray}
また，求める答えは，
$$\sum_{v=0}^{K} \sum_{w=1}^{M} dp[N][v][w]$$
∴この問題を空間計算量$O(N^2 M)$, 時間計算量$O(N^2 M^2)$で解くことが出来ました。
(配列のサイズと型を入力すると何MBか教えてくれるうし)
まぁ間に合う訳がないんですが...
頭を少しだけ働かせると色ごとに重みがない(等価)なことを思い出します。
つまり，色に関する情報は不要なんですね。したがって，次のような$dp$が成り立ちます。
$dp[x][y]:=$先頭$x$個のブロックを隣接かつ同色な組が$y$組になるように塗るような塗り方の総数
遷移は，
\begin{eqnarray} dp[x][y] = \begin{cases} M & ( x = 1 かつ y = 0 ) \\ dp[x-1][y-1] + (M - 1) dp[x-1][y] & ( x > 1 かつ x > y ) \\ 0 & (otherwise) \\ \end{cases} \end{eqnarray}
(真ん中の式の捕捉:いる？)
このとき，求める答えは
$$\sum_{y=0}^{K} dp[N][y]$$
時間，空間計算量ともに$O(N^2)$なのでまだ間に合いません。
ここで，以下の等式(本質)に気づきます。
$$dp[P][y] = dp[P-y][0] \times _{P-y}H_y$$
これは，$P$個のブロックを隣接かつ同色な組が$y$組できるように並べることは，隣接かつ同色な組を全て同一と見做して$P-y$個のブロックを並べ，あとからどの$y$個のブロックを(重複を許して)複製するかを考えても同じなので成立します。
さて，上記の$dp$で求めたいものは
$$\sum_{y=0}^{K} dp[N][y]$$
でした。つまり，$dp[N][y]$の値が分かればよいのですが，上記の等式より，$dp[x][0]$の値が分かれば$dp[N][y]$の値が分かるので，結局，次のような$dp$になります。
$dp[x]:=$先頭$x$個のブロックを隣接かつ同色な組がないように塗るような塗り方の総数
遷移は，
\begin{eqnarray} dp[x][y] = \begin{cases} M & (x = 1) \\ (M - 1) dp[x-1] & ( x > 1) \\ \end{cases} \end{eqnarray}
求める答えは
$$\sum_{y=0}^{K} dp[N-y] \times _{N-y}H_y$$
$$= \sum_{y=0}^{K} M(M - 1)^{N-y-1} \times _{N-y+y-1}C_y$$
$$= M \sum_{y=0}^{K} (M - 1)^{N-y-1} \times _{N-1}C_y$$
これで，時間，空間計算量ともに$O(N)$で解くことが出来ました。
結局，直接数え上げるのとコードは同じになるのですが，こういうアプローチもあるよ～程度に紹介したかったので，しました。(雑に書いた部分はいつかちゃんと書きます。)
