[ARC170C] Prefix Mex Sequence

https://atcoder.jp/contests/arc170/tasks/arc170_c

给定 \(n, m, S_1\sim S_n\)，当 \(S_i = 0\) 时 \(A_i \neq \mathrm{mex}(A_1\sim A_{i - 1})\)，反之 \(A_i = \mathrm{mex}(A_1\sim A_{i - 1})\)，求值域为 \([0, m]\) 的 \(A\) 的数量 \(\bmod\ 998244353\)。

\(1\le n \le 5000, 0\le m\le 10^9\)。

看到题目就会想到 MEX Counting 的刻画方法：延后确立元素。但是这样做非常繁杂，并且 \(\mathcal O(n^3)\) 不存在什么好的优化方法。

尝试优化的过程中我们会产生一种感觉：在状态里记 MEX 似乎完全没用。

看一看这能带来什么启发。对于 \(m\ge n - 1\) 的情况，显然 MEX 就是没有用，但是受制于上面的思路，并没有什么好的进展。

一个 tricky 的点是，关注序列的构造过程。

对于 \(m\ge n -1\)，我们发现 \(A_i\) 仅和 \(A_1\sim A_{i - 1}\) 有关，并且不管 \(S_i\) 为 \(0\) 还是 \(1\)，它们的贡献都可以直接得出。

对于 \(m < n - 1\)，提示已经相当明显，我们只需要记录 \([0, m]\) 中还剩多少数没有填就能继续维持上述的贡献计算方式。

更进一步地，可以发现这样做在这个题可行的原因是，一个元素的 地位，包括其贡献和取值限制，在不同的状态中完全一致，于是不需要关注状态，而只需要关注这个元素本身。

#include <bits/stdc++.h>
#define pb emplace_back
#define fir first
#define sec second

using i64 = long long;
using pii = std::pair<int, int>;

constexpr int maxn = 5e3 + 5;
constexpr int mod = 998244353;

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

int power(int x, int y) {
	int ans = 1;
	for (; y; y >>= 1) {
		if (y & 1) ans = (i64)ans * x % mod;
		x = (i64)x * x % mod;
	}
	return ans;
}

void add(int& x, int y) { if ((x += y) >= mod) x -= mod; return ; }

int main() {
	scanf("%d %d", &n, &m);
	for (int i = 1; i <= n; ++i)
		scanf("%d", &a[i]);
	if (m >= n - 1) {
		int ans = 0;
		for (int i = 1; i <= n; ++i)
			ans += !a[i];
		printf("%d\n", power(m, ans));
		return 0;
	}
	f[0][m + 1] = 1;
	for (int i = 1; i <= n; ++i) {
		if (a[i] == 1) {
			for (int j = 0; j <= m; ++j)
				f[i][j] = f[i - 1][j + 1];
		} else {
			for (int j = 0; j <= m + 1; ++j) {
				if (j >= 1) add(f[i][j], (i64)f[i - 1][j + 1] * j % mod);
				add(f[i][j], (i64)f[i - 1][j] * (m + 1 - j) % mod);
			}
		}
	}
	int ans = 0;
	for (int i = 0; i <= m + 1; ++i)
		add(ans, f[n][i]);
	printf("%d\n", ans);
	return 0;
}