[abc279 G] At Most 2 Colors

G - At Most 2 Colors (atcoder.jp)

重点讲解方法三，因为方法三是蒟蒻都能想出来的方法一和方法二都可以借助方法三的思想推出

方法一

这是最简单的设置状态的方法，\(dp[i]\)表示前\(i\)个的方案数，然后分类

• 若\([i-k+1,i-1]\)有两种颜色

  那么第\(i\)位的取值肯定时这两种颜色中的一个，所以就是要\(\times2\)，考虑如何使得\([i-k+1,i-1]\)必须有两个颜色，用总方案减去只有一种颜色的方案，也就是\(dp[i-1]-dp[i-k+1]\)

\(dp[i]+=(dp[i-1]-dp[i-k+1])\times2\)

• 若\([i-k+1,i-1]\)只有一种颜色

此时第\(i\)位有\(c\)种取值，则\(dp[i]+=dp[i-k+1]\times c\)

所以综上：\(dp[i]=dp[i-1]\times2+(c-2)\times dp[i-k+1]\)

方法二

设\(dp[i][1/2]\)表示\([i-k+2,i]\)中有\(1/2\)种颜色的方案数

分类讨论，然后枚举\(j\)表示\([j+1,i]\)的颜色都相同

\[\begin{aligned}
&dp[i][1]=\sum_{j=1}^{i-k+1}(dp[j][1]\times(c-1)+dp[j][2])\\
&dp[i][2]=\sum_{j=i-k+2}^{i-1}(dp[j][1]\times(c-1)+dp[j][2])
\end{aligned}
\]

方法三

\(dp\)状态同二

\[\begin{aligned}
&dp[i][1]=dp[i-1][1]+dp[i-k+1][1]\times(c-1)+dp[i-k+1][2]\\
&dp[i][2]=dp[i-1][1]\times(c-1)+dp[i-1][2]\times2-dp[i-k+1][1]\times(c-1)
-dp[i-k+1][2]
\end{aligned}
\]

• 对于\(dp[i][1]\)

  • 当\(i-k+1\)的颜色与\([i-k+2,i]\)的颜色相同时，显然有\(dp[i-1][1]\)

  • 颜色不同时，将考虑的范围向前扩展到\([i-2k+3,i-k+1]\)，这时可以保证对于以\([i-k+1,i-1]\)结尾的长度为\(k-1\)的区间都被囊括其中，以下所有方法的合法性都可以用这个来解释

    • 当\([i-2k+3,i-k+1]\)的颜色都相同时，\([i-k+2,i]\)的颜色可以取除了\([i-2k+3,i-k+1]\)以外的任何颜色，这时就是\(dp[i-k+1][1]\times(c-1)\)

    • 当\([i-2k+3,i-k+1]\)颜色不同时，若两种颜色分别为\(a\)和\(b\)，且\(i-k+1\)的颜色为\(a\)，\([i-k+2,i]\)的为\(b\)，这时就是\(dp[i-k+1][2]\)

• 对于\(dp[i][2]\)

  • 当\([i-k+1,i-1]\)的颜色相同时，显然有\(i\)的颜色取值有除了\([i-k+1,i-1]\)的颜色外的所有颜色，也就是\(dp[i-1][1]\times(c-1)\)

  • 当\([i-k+1,i-1]\)的颜色不同时，\(i\)的取值就有两种，这时有\(dp[i-1][2]\times2\)，但这并不合法，因为\(dp[i-1][2]\)中包含了\([i-k+2,i-1]\)的颜色都为\(a\)，而\(i-k+1\)的颜色为\(b\)的方案数，所以考虑减去这种不合法的方案

    发现这种不合法的情况就是\(dp[i][1]\)的第二个大类，所以直接使用

综上：

\[\begin{aligned}
&dp[i][1]=dp[i-1][1]+dp[i-k+1][1]\times(c-1)+dp[i-k+1][2]\\
&dp[i][2]=dp[i-1][1]\times(c-1)+dp[i-1][2]\times2-dp[i-k+1][1]\times(c-1)
-dp[i-k+1][2]
\end{aligned}
\]

#include<bits/stdc++.h>
#define ll long long
using namespace std;
const int N=1e6+5,MOD=998244353;
int n,k,c;
ll dp[N][2],invf=499122177;
int main(){
	scanf("%d%d%d",&n,&k,&c);
	dp[1][1]=c%MOD;
	for(int i=2;i<=n;++i){
		dp[i][1]=((dp[i-1][1]+dp[max(0,i-k+1)][1]*(c-1)%MOD)%MOD+dp[max(0,i-k+1)][2])%MOD;
		dp[i][2]=((dp[i-1][1]*(c-1)%MOD+dp[i-1][2]*2%MOD)%MOD-(dp[max(0,i-k+1)][1]*(c-1)%MOD+dp[max(0,i-k+1)][2])%MOD+MOD)%MOD;
	}
	printf("%lld\n",(dp[n][1]+dp[n][2])%MOD);
	return 0;
}