【数据结构与算法】Z算法（扩展KMP）（C++和Python写法）

news2024/9/30 6:19:51

Z算法（扩展KMP）

文章目录

Z算法（扩展KMP）
- 朴素求法
- 线性求法
- 力扣类型题
- - 变种题：[3303. 第一个几乎相等子字符串的下标](https://leetcode.cn/problems/find-the-occurrence-of-first-almost-equal-substring/)

所谓Z算法，就是求一个字符串中，每个后缀子串和主串的前缀匹配字符数的数组，其也成为Z数组

eg：主串为aaaab(首位总为0，因为包含首位即本体，无意义)

aaaab aaab -> 3
aaaab aab -> 2
aaaab ab -> 1
aaaab b -> 0
结果集[0, 3, 2, 1,0]

朴素求法

时间复杂度为O(n^2)，暴力获取Z数组。

每次都从头匹配，如果符合往后++，不符合则返回，下一次又从头匹配。

vector<int> z_function_trivial_simple(string s)
{
    int n = (int)s.length();
    vector<int> z(n);
    for (int i = 1; i < n; ++i)
    {
        while (i + z[i] < n && s[z[i]] == s[i + z[i]])
            ++z[i];
    }
    return z;
}

线性求法

我们使用一个滑动窗口[l,r]，这个滑动窗口总是往右移动，我们可以称之为Z_box

这个z_box具有特性：s[l, r] = s[0, r-l]（s为字符串，l和r总是从0开始）

我们再次复习一下z数组的含义：z[i]表示从s[i]开始直到末尾的子字符串和s整个字符串匹配的前缀和

问题一：如何获取这个滑动窗口？

由于滑动窗口（z_box）总是向右移动，所以我们要用z数组及i来辅助获取。

具体方法为：当i+z[i] -1 > r时，修改l和r的位置，是l = i , r = i + z[i] - 1

原因：1. 我们希望滑动窗口会比需要匹配的数字更靠后，或者说能够包含未来匹配的位置，并且滑动窗口总是往右的。

i这里代表新窗口的起始位，z[i]代表匹配的长度， -1 是因为z[i]的数字里包含i的位置。

换句话说，所谓新的z_box就是更往右的匹配上的子串前缀。这么说可能比较抽象，请以下图例辅助理解：

问题二：这个滑动窗口的具体作用？

这个滑动窗口只在i ∈[l, r]时发生作用。

我们以上图例作为一个例子，作为讲解：

此时 i = 5 ，5包含在[4,6]中，而且刚好是中间
因为 s[0,2] == s[4,6] ，那么z[5] 可以直接参考z[1]获取

== > 即z[i] = z[i - l]
但这只是上图的可能性，因为上图中z[i-l] == 1 这个值小于r - i + 1 -> 6- 5 + 1 -> 2，我们已经知道了最多只能匹配到这里

但是！还有一种可能，就是z[i-1] == (r - i + 1)，这种情况我们无法预测r后面是否可以继续匹配，那么我就需要从r的后一位开始匹配。而这种匹配方式则回到了原始的匹配中，不再进行讲解，但是这种情况我们依然可以省略已经处于滑动窗口中的匹配。

下面代码展示（如果还不理解：可以用这个网站模拟：演示Z函数）

C++ 代码

vector<int> z_function(string s)
{		
    vector<int> z(s.size(), 0);
    int l = 0, r = 0;
    for (int i = 1; i < s.size(); i++)
    {
        if (i <= r && z[i - l] < r - i + 1)
        {
            z[i] = z[i - l];
        }
        else 
        {
            z[i] = max(0, r - i + 1);
            // 从头开始暴力求解
            while (i + z[i] < s.size() && s[z[i]] == s[i + z[i]])
                ++z[i];
        }
        if (i + z[i] - 1 > r)
        {
            l = i, r = i + z[i] - 1;
        }
        // 可以打印进行看看
        cout << "i: "<< i << ", z[i]: "<< z[i] << ", [l, r]: ["<< l <<", " << r<<"]"<<endl;
    }
    return z;
}

Python代码

def getZArray(self, s : str) -> List[int]:
    # z[i] 为从i开始能和主串从头匹配的字符总数
    z = [0] * len(s)
    l, r = 0, 0
    for i in range(1, len(s)):
        # 当i在窗口内
        # 如果z[i-l] < (r-i+1)，说明z[i-l]能匹配的字符数已经可知，直接获取
        # 否则，有可能超出这个数字，需要从末尾继续暴力寻找
        if i <= r:  # i在窗口内
            z[i] = min(z[i - l], r - i + 1)
        while i + z[i] < len(s) and s[z[i]] == s[i + z[i]]:  # 暴力匹配剩余部分
            z[i] += 1
        if i + z[i] - 1 > r:  # 更新窗口边界
            l, r = i, i + z[i] - 1
    return z

力扣类型题

变种题：3303. 第一个几乎相等子字符串的下标

这道题在Z算法的基础上，变形为前缀+后缀的组合，详情可以看这篇题解，写得很好，我不班门弄斧了。贴上我的代码。

C++

class Solution {
public:
	int minStartingIndex(string s, string pattern) {
		int m = pattern.size(), n = s.size();
		string combine = pattern + s;
		reverse(pattern.begin(), pattern.end());
		reverse(s.begin(), s.end());
		string combinervs = pattern + s;
		vector<int> pre = getZArray(combine);			// pre_l = z[m+l]
		vector<int> suf = getZArray(combinervs);		// suf_r = z[m+(n-r-1)]
		for (int l = 0, r = m - 1; r < n; l++, r++)
		{
			if (pre[m + l] + suf[m + (n - r - 1)] + 1 >= m)
				return l;
		}
		return -1;
	}

private:
	vector<int> getZArray(string& s)
	{
		vector<int> z(s.size(), 0);
		int l = 0, r = 0;
		for (int i = 1; i < s.size(); i++)
		{
			if (i <= r && z[i - l] < r - i + 1)
			{
				z[i] = z[i - l];
			}
			else 
			{
				z[i] = max(0, r - i + 1);
				while (i + z[i] < s.size() && s[z[i]] == s[i + z[i]])
					++z[i];
			}
			if (i + z[i] - 1 > r)
			{
				l = i, r = i + z[i] - 1;
			}
		}
		return z;
	}
};

Python

from typing import List

class Solution:
    def getZArray(self, s: str) -> List[int]:
        # z[i] 是从索引 i 开始的子串与主串前缀匹配的长度
        z = [0] * len(s)
        l, r = 0, 0
        for i in range(1, len(s)):
            if i <= r:  # i在窗口内
                z[i] = min(z[i - l], r - i + 1)
            while i + z[i] < len(s) and s[z[i]] == s[i + z[i]]:  # 暴力匹配剩余部分
                z[i] += 1
            if i + z[i] - 1 > r:  # 更新窗口边界
                l, r = i, i + z[i] - 1
        return z

    def minStartingIndex(self, s: str, pattern: str) -> int:
        m, n = len(pattern), len(s)
        
        # 生成前缀和后缀Z数组
        combined = pattern + s
        reversed_combined = pattern[::-1] + s[::-1]
        pre = self.getZArray(combined)
        suf = self.getZArray(reversed_combined)
        
        # 检查匹配位置
        for l in range(n - m + 1):
            r = l + m - 1
            if pre[m + l] + suf[m + (n - r - 1)] + 1 >= m:
                return l
        return -1