题目
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;
}