G - At Most 2 Colors Editorial by math_hiyoko

listを使ったdpテーブルのシフトによる解法

この解法では

\(dp[i][j] = (マス1〜iの色を塗ったとき、直前ちょうどjマスが同じ色となる場合の数)\)

とします。(ただし\(K-1 \leq i \leq N, 1 \leq j \leq K-1\))
このdpテーブルの遷移の式は以下のようになります(詳しい説明は一番下を見てください)。

\( dp[i + 1][j] = \begin{cases} dp[i][K-1] \times (C - 2) + \sum_{j = 1}^{K-1} dp[i][j] & (j = 1)\\ dp[i][j-1] & (2 \leq j \leq K - 2)\\ dp[i][j-1] + dp[i][K-1] & (j = K - 1) \end{cases} \)

この遷移を愚直に行うと\(O(NK)\)となりますが、リストを使ってdpテーブルを使い回すと\(O(N)\)にできます。
具体的には、\( 2 \leq j \leq K - 2\)のときに\(dp[i+1][j]\)\(dp[i][j-1]\)が等しいことを利用し、\(dp[i]\)に対して末尾の要素の削除と先頭への要素の追加を1回ずつ行うことで\(O(1)\)\(dp[i+1]\)に変更します。
C++ではstd::listやstd::dequeを使ってこれを実装することができます。
実装例

dpテーブルの遷移の説明

貰うdpの形で考えます。

\(dp[i + 1][1] \)への遷移

この場合、マス\(i\)とマス\(i+1\)に異なる色が塗られていることになります。
・マス\(i\)からみて直前\(K-1\)が全て同じ色であればマス\(i+1\)に塗ることができるのは\(C-1\)色 ・それ以外であればマス\(i+1\)に塗ることができるのは\(1\)色 なので、
\(dp[i+1][1]\\= dp[i][K-1] \times (C - 1) + \sum_{j = 1}^{K-2} dp[i][j]\\= dp[i][K-1] \times (C - 2) + \sum_{j = 1}^{K-1} dp[i][j]\)
となります。

\(dp[i + 1][j](2 \leq j \leq K - 2)\) への遷移

この場合、マス\(i\)とマス\(i+1\)には同じ色が塗られていることになるので、
\(dp[i+1][j] = dp[i][j-1]\)
となります。

\(dp[i + 1][K-1]\) への遷移

この場合もマス\(i\)とマス\(i+1\)には同じ色が塗られていますが、\(dp[i][K-1]\)も寄与するので
\(dp[i+1][K-1]=dp[i][K-2]+dp[i][K-1]\)
となります。

posted:
last update: