[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

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

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

原文链接: https://www.cnblogs.com/mekoszc/p/17032436.html

欢迎关注

微信关注下方公众号，第一时间获取干货硬货；公众号内回复【pdf】免费获取数百本计算机经典书籍

原创文章受到原创版权保护。转载请注明出处：https://www.ccppcoding.com/archives/310313

非原创文章文中已经注明原地址，如有侵权，联系删除

关注公众号【高性能架构探索】，第一时间获取最新文章

转载文章受原作者版权保护。转载请注明原作者出处！