【Codeforces 1093F】Vasya and Array / Solution

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.

  1. $i≥len$, as we should have at least $len$ elements;
  2. 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]);
}
Show Comments