Description
Vasya has got an array consisting of $n$ integers, and two integers $k$ and $len$ in addition. All numbers in the array are either between $1$ and $k$ (inclusive), or equal to $−1$. The array is good if there is no segment of $len$ consecutive equal numbers.
Vasya will replace each $−1$ with some number from $1$ to $k$ (inclusive) in such a way that the resulting array is good. Tell him the number of ways to do this replacement. Since the answer may be large, print it modulo $998244353$.
Idea
Let’s try dynamic programming approach to this problem. Let $f_{i,j}$ be the number of ways to replace all $−1$ with numbers from $1$ to $k$ in such a way that array $\{a_1,\cdots,a_i\}$ is good and the last number of that array is $j$, and $s_i$ be $\sum f_{i,j}$.
Then initially it’s $f_{i,j}=s_{i-1}$ if $a_i$ equals to −1−1 or $j$. However, we could include incorrect states doing this directly. To subtract all the bad states, let’s talk about when it would happen.
- $i≥len$, as we should have at least $len$ elements;
- segment $\{a_1,\cdots,a_i\}$ has all its elements either equal to $−1$ or $j$.
If both of these conditions hold then you should subtract all the bad states from $f_{i,j}$, which equals to $s_{i-len}-f_{i-len,j}$.
Code
#include <bits/stdc++.h>
#define md 998244353
#define maxn 100001
#define max(a,b) (a>b?a:b)
#define maxk 101
int n,k,len,a[maxn];
int f[maxn][maxk],s[maxn],cnt[maxn][maxk];
void inc(int &a,int b){a=((a+b>=md)?a+b-md:a+b);}
int main(){
scanf("%d%d%d",&n,&k,&len);
for (register int i=1;i<=n;++i)
scanf("%d",&a[i]);
for (register int i=1;i<=n;++i)
for (register int j=1;j<=k;++j)
inc(cnt[i][j],cnt[i-1][j]+(a[i]==j||a[i]==-1));
for (register int i=1;i<=n;++i){
for (register int j=1;j<=k;++j){
if (!(a[i]==j||a[i]==-1)) continue;
int add=1;
if (i>1) add=s[i-1];
inc(f[i][j],add);
bool ok=i>=len;
int l=max(1,i-len+1);
ok&=(cnt[i][j]-cnt[l-1][j]==len);
if (!ok) continue;
if (i==len) {inc(f[i][j],md-1);continue;}
int sum=f[i-len][j];
inc(sum,md-s[i-len]);
inc(f[i][j],sum);
}
for (register int j=1;j<=k;++j)
inc(s[i],f[i][j]);
}
printf("%d\n",s[n]);
}