Problem Statement
There is a grid with $1 \times N$ squares, numbered $1,2,\dots,N$ from left to right.
Takahashi prepared paints of $C$ colors and painted each square in one of the $C$ colors.
Then, there were at most two colors in any consecutive $K$ squares.
Formally, for every integer $i$ such that $1 \le i \le N-K+1$, there were at most two colors in squares $i,i+1,\dots,i+K-1$.
In how many ways could Takahashi paint the squares?
Since this number can be enormous, find it modulo $998244353$.
Constraints
- All values in the input are integers.
- $2 \le K \le N \le 10^6$
- $1 \le C \le 10^9$
Input
The input is given from Standard Input in the following format:
$N$ $K$ $C$
Output
Print the answer as an integer.
Sample Input 1
3 3 3
Sample Output 1
21
In this input, we have a $1 \times 3$ grid.
Among the $27$ ways to paint the squares, there are $6$ ways to paint all squares in different colors, and the remaining $21$ ways are such that there are at most two colors in any consecutive three squares.
Sample Input 2
10 5 2
Sample Output 2
1024
Since $C=2$, no matter how the squares are painted, there are at most two colors in any consecutive $K$ squares.
Sample Input 3
998 244 353
Sample Output 3
952364159
Print the number modulo $998244353$.
定义 \(dp_{i,j}\) 为前 \(i\) 个位置，满足 \([j,i]\) 颜色全部相等，但是 \(j-1\) 的颜色不等的方案数。
为什么会想到这个定义呢？因为我们会发现，当且仅当倒数 \(k-1\) 个数是相等的，那么我们在下一个位置才可以随便选数字。否则的话，下一个位置如果还要和上一个位置不一样，只有一种选择。
那么对于 \(i>j\)，所有的 \(dp_{i,j}\) 都是一样的，因为往后只能填和前一位用的同一个颜色。\(dp_i\) 表示的是前 \(i\) 位有多少种填法。
然后如果结尾的相同的颜色个数大于等于 \(k-1\)，那么有 \(c-1\) 种填法。否否则，就只有一种填法。这里可以拿前缀和优化。
还有一个特殊情况，就是把 \([2,i-1]\) 全部填成一种颜色，此时有 \(c*(c-1)\) 种情况
#include<bits/stdc++.h>
using namespace std;
const int N=1e6+5,P=998244353;
int n,k,c,s[N];
int main()
{
scanf("%d%d%d",&n,&k,&c);
for(int i=2,dp;i<=n;i++)
{
dp=(s[i-1]+(c-2LL)*s[max(i-k+1,0)]%P+P+c*(c-1LL))%P;;
s[i]=(s[i-1]+dp)%P;
// printf("%d ",dp);
}
printf("%d",(s[n]+c)%P);
}