字符串从入门到退竞（2）——KMP 算法

约定在文字描述中字符串下标从 \(1\) 开始，代码中从 \(0\) 开始。

前缀函数

对于长 \(n\) 的字符串 \(S\)，其前缀函数是一个长 \(n\) 的数组（数列）\(\pi\)，定义如下：

• 若 \(S[1..i]\) 有相等的真前缀和真后缀（称之为 border），\(\pi[i]\) 为其长度的最大值；
• 若不存在，\(\pi[i]=0\)。

字符串的真前缀或真后缀不包括其本身。

例如 \(\mathtt{aabaaab}\) 的前缀函数是 \([0,1,0,1,2,2,3]\)。

计算前缀函数的朴素算法

容易想到对于每个位置进行暴力枚举。

vector<int> prefix_function(const string &s) {
  int n = s.length();
  vector<int> pi(n);
  for (int i = 1; i < n; ++i)
    for (int j = i; j >= 0; --j)
      if (s.substr(0, j) == s.substr(i - j + 1, j)) {
        pi[i] = j;
        break;
      }
  return pi;
}

字符串比较的复杂度是 \(O(n)\)，总时间复杂度为 \(O(n^3)\)。

第一个优化

注意到前缀函数后一个值最多比前一个值多 \(1\)，我们可以对算法改进。

vector<int> prefix_function(const string &s) {
  int n = s.length();
  vector<int> pi(n);
  for (int i = 1; i < n; ++i)
    for (int j = pi[i - 1] + 1; j >= 0; --j)
      if (s.substr(0, j) == s.substr(i - j + 1, j)) {
        pi[i] = j;
        break;
      }
  return pi;
}

这看似是常数优化，实则不然。

应用势能分析，定义势能函数为 \(\Phi(i)=\pi[i]\times n\)，易得每次求新的 \(\pi[i]\) 的均摊复杂度为 \(O(n)\)。

第二个优化

延续第一个优化的思路，我们只需检查字符串的每一个 border 是否可以“扩展”。

字符串 \(S[1..i]\) 最大的 border 长度为 \(\pi[i]\)，容易注意到次长的 border 即最长的 border 的 border，其长度为 \(\pi[\pi[i]]\)，以此类推。

vector<int> prefix_function(string s) {
  int n = s.length();
  vector<int> pi(n);
  for (int i = 1; i < n; ++i) {
    int j = pi[i - 1];
    while (j > 0 && s[i] != s[j]) j = pi[j - 1];
    if (s[i] == s[j]) ++j;
    pi[i] = j;
  }
  return pi;
}

时间复杂度为 \(O(n)\)。

这是一个在线算法，在读取数据时立即处理。

KMP 算法

给定文本串 \(T\) 和模式串 \(S\)，我们尝试找到所有 \(S\) 在 \(T\) 中的所有出现。

我们构造字符串 \(S+\mathtt{\#}+T\)，其中 \(\mathtt{\#}\) 是既不出现在 \(S\) 也不出现在 \(T\) 中的分割符。计算该字符串的前缀函数，显然不会出现 \(\pi[i]>|S|\)；若 \(\pi[i]=|S|\) 成立，意味着 \(S\) 在这里完整出现了一次。

vector<int> find_occurrences(string text, string pattern) {
  string cur = pattern + '#' + text;
  int sz1 = text.size(), sz2 = pattern.size();
  vector<int> v;
  vector<int> lps = prefix_function(cur);
  for (int i = sz2 + 1; i <= sz1 + sz2; i++) {
    if (lps[i] == sz2) v.push_back(i - 2 * sz2);
  }
  return v;
}

字符串的周期

若字符串 \(S\) 有长 \(r\) 的 border，则 \(|S|-r\) 显然是 \(S\) 的周期。求前缀函数后枚举 border 即可。

统计每个前缀的出现次数

给定长 \(n\) 的字符串 \(S\)，考虑位置 \(i\) 的前缀函数值 \(\pi[i]\)，这意味字符串 \(S\) 长 \(\pi[i]\) 的前缀在 \(i\) 出现并以 \(i\) 作为右端点，并且不存在更长的前缀。同时，更短的前缀可能在这里作为右端点，因此还需统计 \(\pi[\pi[i]],\pi[\pi[\pi[i]]]\)，等等。

vector<int> ans(n + 1);
for (int i = 0; i < n; ++i) ++ans[pi[i]];
for (int i = n - 1; i > 0; --i) ans[pi[i - 1]] += ans[i];
for (int i = 0; i <= n; ++i) ++ans[i];

前缀自动机

在 KMP 算法中，前缀函数不会超过模式串 \(S\) 的长度，因此 \(T\) 部分每个字符的前缀函数只与模式串的前缀函数和前一个字符的前缀函数有关。我们将当前状态定义为前一个字符的前缀函数，可以求出自动机的转移表 \((\pi_\text{old},c)\to \pi_{\text{new}}\)：

void compute_automaton(string s, vector<vector<int>>& aut) {
  s += '#';
  int n = s.size();
  vector<int> pi = prefix_function(s);
  aut.assign(n, vector<int>(26));
  for (int i = 0; i < n; ++i) {
    for (int c = 0; c < 26; ++c) {
      if (i > 0 && 'a' + c != s[i])
        aut[i][c] = aut[pi[i - 1]][c];
      else
        aut[i][c] = i + ('a' + c == s[i]);
    }
  }
}

时间复杂度可以做到 \(O(|\Sigma|n)\)。