Z 算法/拓展 KMP 详解

更好的阅读体验

引入

先来看几个问题：

什么是 Z 算法呢？

Z 算法（Z-Algorithm，又叫拓展 KMP 或 exKMP）是一种字符串匹配算法，可以在 $O(n)$ 时间内算出一些关于字符串匹配的问题。

为什么要用 Z 算法？

暴力进行匹配是 $O(n^2)$ 的，但用 Z 后是 $O(n)$ 的（KMP 也可以实现类似的功能）。

接下来是一些约定：

1. $s[l..r]$ 表示 $s_l{s}_{l+1}{s}_{l+2}...s_r$
2. z 函数用 $z()$ 表示，原字符串用 $s$ 表示
3. 下标从 $1$ 开始

算法流程

定义

定义： 对于任意 $i(2\le i\le n)$，定义 $z(i)$ 表示满足 $s[1..x]=s[i..i+x-1]$ 的最大的 $x$。

例如：
$s=\mathtt{\text{aabbaabb}}$，$z(5)=4$，因为 $s[1..4]=s[5..8]$
$s=\mathtt{\text{aabab}}$，$z(3)=0$
$s=\mathtt{\text{ababa}}$，$z(3)=3$

并且，对于任意 $i(2\le i\le n)$，我们把区间 $[i,i+z(i)-1]$ 叫做一个 Z-Box，共有 $n$ 个 Z-Box。
这个概念有什么用呢？一会求 Z 函数时会用到。

如何求解

求 $z[1..n]$ 的总复杂度是 $O(n)$ 的。

我们知道暴力求 $z$ 是 $O(n^2)$的，那么为了优化复杂度，我们需要用之前算出的 $z[2..i-1]$ 来算出 $z(i)$。

我们维护右端点最右的 Z-Box 的左右端点，记为 $zl$, $zr$。（初始都为 $0$）

考虑如何求 $z(i)$。如果 $zr<i$，那么我们并无法利用之前的 $z$，所以直接暴力更新 $z(i)$，并更新 $zl$, $zr$。下面讨论 $zr\ge i$ 的情况。

令 $i^{\prime }=i-zl+1$，图中绿色是 $i^{\prime }$ 的 Z-Box 和 $s$ 与之长度相等的前缀，由 $z$ 函数定义得两段绿色相等。蓝色同理。两段红色是 $[zl,zr]$ 和 $s$ 与之长度相等的前缀，两段红色也相等。

由图可知，$z(i)>=minz(i^{\prime }),zr-i+1$。

那么如果 $z(i^{\prime })<zr-i+1$，那么直接让 $z(i)=z(i^{\prime })$。

但是如果 $z(i^{\prime })\ge zr-i+1$ 呢？是不是让 $z(i)=zr-i+1$ 就好了呢？

其实不一定，看下面这张图：

图中 $zr-i+1$ 应该等于蓝色，但是实际上 $z(i)$ 应该是蓝色+橙色，所以这种情况我们直接在 $zr-i+1$ 的基础上暴力更新 $z(i)$，并更新 $zl,zr$。

无关紧要的证明：

其实 $z(i^{\prime })$ 不可能大于 $zr-i+1$，因为如果大于，那么就会变成下面这样：

那么粉色和橙色的前面部分是相等的，那么红色就可以继续往后延伸，与前面 $z$ 的定义矛盾。

时间复杂度

因为 $zl$, $zr$ 一直往右移动，不移动 $zl,zr$ 时的操作为常数，所以时间复杂度为 $O(n)$。

代码

int z[N];
void z_algorithm(const char *s) { // 下标从 1 开始
	int l = 0, r = 0;
	int n = strlen(s + 1);
	for(int i = 2; i <= n; i++) {
		if(i <= r) z[i] = std::min(z[i - l + 1], r - i + 1);
		while(i + z[i] <= n && s[i + z[i]] == s[z[i] + 1]) z[i]++;
		if(i + z[i] > r) l = i, r = i + z[i] - 1;
	}
}

应用

例题 1

给出两个字符串 s, t，求出所有的位置 x，满足 $s[x..x+len(t)-1]=t$。
其中 $len(s),len(t)\le {10}^5$，s, t 仅包含大小写字母。

可以设一个字符 $r=t+\mathtt{\text{!}}+s$，然后对 $r$ 跑一边 Z-Algorithm。
然后看 $r[len(t)+2..len(r)]$ 中有多少个的 Z 函数等于 $len(t)$。

也可以按照 Z 算法类似的方式直接算。（设 $e(i)$ 表示满足 $s[i..i+x-1]=t[1..x]$ 的最大的 $x$）

int e[N];
void exkmp(const char *s, const char *t) { // 下标从 1 开始
	int l = 0, r = 0;
	int n = strlen(s + 1);
	for(int i = 2; i <= n; i++) {
		if(i <= r) e[i] = std::min(z[i - l + 1], r - i + 1);
		while(i + e[i] <= n && s[i + e[i]] == t[e[i] + 1]) e[i]++;
		if(i + e[i] > r) l = i, r = i + e[i] - 1;
	}
}

例题 2

洛谷 P5410

#include <cstdio>
#include <algorithm>
#include <cstring>

typedef long long LL;

const int N = 2e7 + 5;

char a[N], b[N];

int z[N];
void z_algorithm(const char *s) { // 下标从 1 开始
	int l = 0, r = 0;
	int n = strlen(s + 1);
	for(int i = 2; i <= n; i++) {
		if(i <= r) z[i] = std::min(z[i - l + 1], r - i + 1);
		while(i + z[i] <= n && s[i + z[i]] == s[z[i] + 1]) z[i]++;
		if(i + z[i] > r) l = i, r = i + z[i] - 1;
	}
}

int e[N];
void exkmp(const char *s, const char *t) { // 下标从 1 开始
	int l = 0, r = 0;
	int n = strlen(s + 1);
	for(int i = 2; i <= n; i++) {
		if(i <= r) e[i] = std::min(z[i - l + 1], r - i + 1);
		while(i + e[i] <= n && s[i + e[i]] == t[e[i] + 1]) e[i]++;
		if(i + e[i] > r) l = i, r = i + e[i] - 1;
	}
}

int main() {
	scanf("%s%s", a + 1, b + 1);
	int la = strlen(a + 1), lb = strlen(b + 1);
	z_algorithm(b);
	z[1] = lb;
	LL ans = 0;
	for(int i = 1; i <= lb; i++) ans ^= (LL)i * (z[i] + 1);
	printf("%lld\n", ans);
	exkmp(a, b);
	for(int i = 1; i <= lb; i++) if(a[i] == b[i]) e[1] = i; else break;
	ans = 0;
	for(int i = 1; i <= la; i++) ans ^= (LL)(i) * (e[i] + 1);
	printf("%lld\n", ans);
	return 0;
}