問題概要

問題文はこちらを参照してください。

$O(nk^2)$の解法までしか自分ではわかりませんでした。

解法

制限から$O(n)$の解法でなければならないというのはわかるのですが、どのようにそこにたどりつくのかを考えてみます。

駅$i$までで、駅$i$を含んで連続するちょうど$j$個の駅に止まる組合せを考えると、$j=0$(駅$i$に止まらない)の時は駅$i-1$までに連続して$0～K-1$止まる組合せの和になります。$j\gt 0$の時(駅$i$に止まる)は、駅$i-1$までに連続してちょうど$j-1$止まる組合せと同じになります。

漸化式の定義： $$dp[i][j]= \begin{cases}\displaystyle \sum_{k=0}^{K-1}dp[i-1][k] & (j=0)\\ dp[i-1][j-1] & (j\gt 0) \end{cases}$$

$dp[i][j]$：駅$i$までで、駅$i$を含んで連続する$j$個の駅に止まる組合せ数
$(0\le i \le N,\ 0\le j \le K-1)$

初期値：$dp[1][0]=0$(駅1には必ず止まるので),$dp[1][1]=1(駅0には止まらないという仮定なので)$

求める解：$\displaystyle \sum_{k=1}^{K-1}dp[N][k]$(駅$N$には必ず止まるので$k=0$は含まない)

こう考えると、求める解は$dp[N][0]=0$(駅$N$には必ず止まるので)とすると$dp[N+1][0]=\displaystyle \sum_{k=0}^{K-1}dp[N][k]$と同じということがわかります。

次に、$j\gt 0$のケースの式から、例えば$j=K-1$の時、$dp[i][K-1]=dp[i-1][K-2]=...=dp[i-(K-1)][0]$となって、「駅$i$までで、駅$i$を含んで連続する$K-1$個の駅に止まる組合せ数」は、「駅$i-(K-1)$までで、駅$i-(K-1)$に止まらない組合せ数」と同じという式が得られます。確かに、駅$i-(K-1)$より後連続して駅$i$まで止まる組合せは1通りしかないので、そうなります。これらから、結局$j=0$の時だけ考えればよいということになります。

$j=0$の時の漸化式は結局、$dp[i][0]=\displaystyle \sum_{k=0}^{K-1}dp[i-1-k][0]$となり、駅$i$に止まらない場合の組合せ数を考えればよいということになりました。

さらにさらに、1つ前までの$K$個分の和を求めればいいので、累積和で求めることができます。

漸化式の定義： $$\begin{eqnarray} dp[i] & = & \sum_{k=0}^{K-1}dp[i-1-k]\\ & = & \sum_{j=0}^{i-1}dp[j] - \sum_{j=0}^{i-1-K}dp[j]\\ & = & sum[i-1]-sum[i-1-K]\\ sum[i] & = & \sum_{j=0}^{i}dp[j] \end{eqnarray}$$

$dp[i]$：駅$i$までで、駅$i$に止まらず、$K$個以上連続して止まらない組合せ数
$(0\le i \le N+1)$

初期値：$dp[1]=0,\ dp[N]=0,\ dp[0]=1,\ sum[i]=0(i\lt 0)$

求める解：$dp[N+1]$

考察

あとから考えると、このような説明を付けることはできるのですが、本番ではどのように考えるのがいいのでしょうかね。しかし、いきなり$O(n)$のＤＰを考えるのはとても無理な気がしているので、わかるところから考えて少しずつ改善していくのが常道なのでしょうか。

あと、実装に関してですが、引き算のmodを取るところがあるので、注意が必要です。