[ABC279G] At Most 2 Colors

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);
}