AtCoder Beginner Contest 279 G. At Most 2 Colors(计数/组合数学/dp递推)

题目

n(2<=n<=1e6)个格子，从左到右一字排开，

现在需要给格子涂色，有c(1<=c<=1e9)种颜色，

要求连续的k(2<=k<=n)个格子的颜色数最多只有2种，

求方案数，答案对998244353取模

思路来源

TOYOTA SYSTEMS Programming Contest 2022(AtCoder Beginner Contest 279)（G.At Most 2 Colors容斥,DP,组合数学） - 知乎

题解

dp[i]表示长度为i的格子的合法涂色数，考虑第i个怎么放

①第i个前面k-1个位置有2种颜色，则第i个位置只能放这两种颜色中的一种

用合法方案减只有一种的方法，即得两种颜色的方案数

而只有一种颜色的方案数，等于f[i-k+1]，

此时，让中间的k-2个格子都和i-k+1涂一样的即可

②第i个前面k-1个位置仅有一种颜色，刚求过，f[i-k+1]，

此时，第i个位置有c种选择可以涂，

dp[i]=(dp[i]-dp[i-k+1])*2+c*dp[i-k+1]

当i-k+1<1时，相当于前缀没有限制，

而目的是为了让这中间夹着的<=k-2个格子同色，

此时方案数有c种，与dp[1]统一

注意到c=1或2时，c-2<=0，可以特判，

但代入式子后发现，可以和式子统一

代码

#include<bits/stdc++.h>
using namespace std;
const int N=1e6+10,mod=998244353;
int n,k,c,dp[N];
int main(){
    cin>>n>>k>>c;
    dp[1]=c;
    for(int i=2;i<=n;++i){
        dp[i]=(2ll*dp[i-1]%mod+1ll*(c-2)*dp[max(1,i-k+1)]%mod)%mod;
    }
    cout<<dp[n]<<endl;
	return 0;
}