【洛谷 P4548 [CTSC2006]歌唱王国】【概率生成函数+KMP】

题意

给一个长度为 $n$ 的序列 $A$。有一个空序列 $B$，每次会等概率随机往该序列的末尾加入 $1$ 到 $m$ 中的一个整数，若 $A$ 成为了 $B$ 的子串，则停止。求序列 $B$ 的期望长度。
$n\le 10^5$

分析

定义一个离散随机变量 $X$ 的概率生成函数为$$F(z)=\sum _{i\ge 0}P(X=i)z^i$$

其中 $P(X=i)$ 表示 $X$ 取值为 $i$ 的概率。不难得到 $F(1)=1,F^{\prime }(1)=E(X)$。

回到题目，设 $a_i$ 表示 $A$ 长度为 $i$ 的前缀是否是 $A$ 的一个 border，是则为 $1$，否则为 $0$。设 $f_i$ 表示结束时 $B$ 的长度恰好为 $i$ 的概率，其概率生成函数为 $F(x)$。$g_i$ 表示 $B$ 的长度为 $i$ 时还没结束的概率，生成函数为 $G(x)$。

则有等式$$F(x)+G(x)=1+xG(x)\text{\hspace{1em}(1)}$$

$$G(x)(\frac{1}{m}x{)}^n=\sum _{i=1}^n{a}_{i}F(x)(\frac{1}{m}x{)}^{n-i}\text{\hspace{1em}(2)}$$

$(1)$ 式可以理解为，在序列末尾随机加入一个数字后，有可能结束，也有可能不结束。

$(2)$ 式可以理解为，在一个还未结束的串后面加入原串，则一定会结束。但可能还没全部加入就已经结束。右边的 $i$ 枚举的是结束后还没被加入部分的长度，显然这部分一定是原串的 border。

对 $(1)$ 式求导并代入 $x=1$，可得$$F^{\prime }(x)+G^{\prime }(x)=G(x)+x{G}^{\prime }(x)$$

$$F^{\prime }(1)=G(1)$$

对 $(2)$ 式代入 $x=1$ 可得$$G(1)=\sum _{i=1}^n{a}_i{m}^{i}F(1)=\sum _{i=1}^n{a}_i{m}^i$$

用 KMP 求出 $a_i$ 即可。时间复杂度 $O(n)$。

代码

#include<bits/stdc++.h>
using namespace std;

const int N = 100005;
const int mo = 1e4;

int n, t, a[N], nx[N], m;
bool vis[N];

void get_next(int n)
{
	nx[0] = -1; nx[1] = 0;
	int i = 2, j = 0;
	while (i <= n)
		if (j == -1 || a[j + 1] == a[i])
		{
			nx[i] = ++j;
			i++;
		}
		else j = nx[j];
}

void pri(int x)
{
	vector<int> vec;
	while (x) vec.push_back(x % 10), x /= 10;
	while (vec.size() < 4) vec.push_back(0);
	for (int i = 3; i >= 0; i--) putchar(vec[i] + '0');
	puts("");
}

int main()
{
	scanf("%d%d", &n, &t);
	while (t--)
	{
		scanf("%d", &m);
		for (int i = 1; i <= m; i++) scanf("%d", &a[i]), vis[i] = 0;
		get_next(m);
		int ans = 0, k = m;
		while (k) vis[k] = 1, k = nx[k];
		for (int i = 1, w = n % mo; i <= m; i++, w = w * n % mo) (ans += vis[i] * w) %= mo;
		pri(ans);
	}
	return 0;
}