[题解] [CTSC2006] 歌唱王国

给定 $n$ ，代表有 $n$ 种字符。再给出多个数组 $a$ ，记其长度为 $m$，$1\le a_i\le n$。每次随机写下出一个字符，求第一次写下这个数组（即写下的字符串后缀为该数组）期望要写多少个字符。

$1\le n,m\le 10^5$

题目链接

---

思路

考场上遇到的题稍微有点不一样，是让你对每个前缀都求出其第一次被写下的时间，并且数据范围比较小。当时 yy 了一个理论最劣 $O(n\cdot m^2)$ 的做法，当时还以为是 $O(n\cdot m)$ 的（草），然后莫名奇妙就卡过去了。考完后神仙说是论文题，$O(m)$ 的。然后想了想发现自己的做法是可以做到 $O(m)$ 的。

啊，标准做法可以去看洛谷一大片生成函数详解，甚至知乎上也有（？）。这里提供一个小白做法。

我们令 $f_i$ 为写下前 $i$ 个字符期望要多少步，然后可以写出转移方程：
$$
f_{i+1}=f_i+1+\frac{1}{n}\sum_{c\ne s_{i+1}}f_{i+1}-f_{tar(i,c)}
$$
也就是说，在写下前 $i$ 个字符的情况下，首先需要再写一个字符，如果这个字符恰好为 $s_{i+1}$ ，那么就没有额外贡献了，否则的话会失配跳到 $tar(i, c)$ 。

然后我们移一下项，乘一个 $n$：
$$
f_{i+1}=n(f_i+1)-\sum_{c\ne s_{i+1}}f_{tar(i, c)}
$$
注意到 $tar(i, c)\le i$ ，因此我们就可以依次递推求得 $f$ 了。时间复杂度 $O(n\cdot m)$。

考虑优化。我们发现每次枚举 $c$ 是我们算法的时间复杂度瓶颈。我们记 $\sum_c f_{tar(i, c)}$ 为 $sum_i$，然后动态维护这个 $sum$ 。我们发现，每次新扫到一个位置 $i$ ，只且只会对 $sum_{i-1}$ 造成影响（ $tar(i-1, c)$ 从 $pre_i$ 变为了 $i$ ，$pre$ 代表 kmp 中所求的所有前缀的 border ）。于是我们每次先把 $f_{pre_i}$ 从 $sum_{i-1}$ 中减去，然后再加上 $f_i$ 即可。于是得出下面的代码：

const int mod_n = n % mod; // 注意 n 可能大于 mod
sum[0] = 0;
for (int i = 1; i <= m; ++i) {
	Sub(sum[i - 1], f[pre[i]]); // 首先减去
	f[i] = (qe)f[i - 1] * n % mod;
	Sub(f[i], sum[i - 1]); // 递推式中要求 c != s[i + 1]，这里正好
	Add(f[i], mod_n); // 算完 f[i] 了

	Add(sum[i - 1], f[i]); // 加上
	sum[i] = sum[pre[i]]; // 继承自己 pre 的 sum
}

然后就是 $O(m)$ 的复杂度了。所求答案即为 $f_m$，然后每个前缀对应的答案也分别对应到 $f_i$。

代码

完整代码（mivik.h）：

// Mivik 2020.9.18
#include <mivik.h>

#include <iomanip>

MI cin;

typedef long long qe;

const int $m = 100001;
const int mod = 10000;

inline int add(int x, int y) { return (x += y) >= mod? x - mod: x; }
inline void Add(int &x, int y) { if ((x += y) >= mod) x -= mod; }
inline int sub(int x, int y) { return (x -= y) < 0? x + mod: x; }
inline void Sub(int &x, int y) { if ((x -= y) < 0) x += mod; }

int n, m;
int a[$m], pre[$m], f[$m], sum[$m];
inline void work() {
	cin > m;
	for (int i = 1, j = 0; i <= m; ++i) {
		cin > a[i];
		while (j && a[i] != a[j + 1]) j = pre[j];
		if (i != j + 1 && a[i] == a[j + 1]) ++j;
		pre[i] = j;
	}
	const int mod_n = n % mod;
	sum[0] = 0;
	for (int i = 1; i <= m; ++i) {
		Sub(sum[i - 1], f[pre[i]]);
		f[i] = (qe)f[i - 1] * n % mod;
		Sub(f[i], sum[i - 1]);
		Add(f[i], mod_n);

		Add(sum[i - 1], f[i]);
		sum[i] = sum[pre[i]];
	}
	cout < std::setw(4) < std::setfill('0') < f[m] < endl;
}
int main() {
	cin > n;
	for (int T = R; T; --T) work();
}