AtCoder Regular Contest 170 C Prefix Mex Sequence

洛谷传送门

AtCoder 传送门

我们实际上并不关心 \(\text{mex}\) 的具体值，只关心它有没有成为 \(\text{mex}\)。

考虑有一个 \(k = \min(m + 1, n)\) 个空位的长条。我们每次可以往长条最左边的空位放一个球（对应 \(a_i\) 成为 \(\text{mex}\)），或者往长条不是最左边的空位或者非空位放一个球（对应 \(a_i < k\) 且 \(a_i\) 不是 \(\text{mex}\)），或者直接把这个球扔掉（对应 \(a_i \ge k\)）。

那么直接设 \(f_{i, j}\) 表示考虑了前 \(i\) 个数，长条占了的空位数为 \(j\) 的方案数。转移考虑 \(a_i\) 的取值个数即可。也就是：

• \(s_i = 1\)：往长条最左边的空位放一个球，有 \(f_{i, j + 1} \gets f_{i - 1, j}\)；
• \(s_i = 0\)：考虑放的球是否占用了空位，有 \(f_{i, j} \gets (j + m - k + 1) f_{i - 1, j}, f_{i, j + 1} \gets (k - j - 1) f_{i - 1, j}\)。

时间复杂度 \(O(n \min(n, m))\)。

code

// Problem: C - Prefix Mex Sequence
// Contest: AtCoder - AtCoder Regular Contest 170
// URL: https://atcoder.jp/contests/arc170/tasks/arc170_c
// Memory Limit: 1024 MB
// Time Limit: 2000 ms
// 
// Powered by CP Editor (https://cpeditor.org)

#include <bits/stdc++.h>
#define pb emplace_back
#define fst first
#define scd second
#define mkp make_pair
#define mems(a, x) memset((a), (x), sizeof(a))

using namespace std;
typedef long long ll;
typedef double db;
typedef unsigned long long ull;
typedef long double ldb;
typedef pair<ll, ll> pii;

const int maxn = 5050;
const ll mod = 998244353;

ll n, m, a[maxn], f[maxn][maxn];

void solve() {
	scanf("%lld%lld", &n, &m);
	for (int i = 1; i <= n; ++i) {
		scanf("%lld", &a[i]);
	}
	ll up = min(m + 1, n);
	f[0][0] = 1;
	for (int i = 1; i <= n; ++i) {
		for (int j = 0; j <= up; ++j) {
			if (a[i]) {
				f[i][j + 1] = (f[i][j + 1] + f[i - 1][j]) % mod;
			} else {
				f[i][j] = (f[i][j] + f[i - 1][j] * (j + m - up + 1)) % mod;
				f[i][j + 1] = (f[i][j + 1] + f[i - 1][j] * (up - j - 1)) % mod;
			}
		}
	}
	ll ans = 0;
	for (int i = 0; i <= up; ++i) {
		ans = (ans + f[n][i]) % mod;
	}
	printf("%lld\n", ans);
}

int main() {
	int T = 1;
	// scanf("%d", &T);
	while (T--) {
		solve();
	}
	return 0;
}