Z 函数（扩展 KMP）

• OI Wiki: https://oi-wiki.org/string/z-func/

题目	难度
2223. 构造字符串的总得分和	困难	模板题
3031. 将单词恢复初始状态所需的最短时间 II	困难	模板题
3045. 统计前后缀下标对 II	困难	Z 函数 + 字典树

定义

对于个长度为 n 的字符串 s。定义函数 z[i] 表示 s 和 s[i,n-1] （即以 s[i] 开头的后缀）的最长公共前缀（LCP）的长度。 z 被称为 s 的 Z 函数。特别地，z[0] = 0。

国外一般将计算该数组的算法称为 Z Algorithm，而国内则称其为 扩展 KMP。

时间复杂度：O(n)。

    private int[] z_function(int n, char[] s) {
        int[] z = new int[n];
        for (int i = 1, l = 0, r = 0; i < n; ++i) {
            if (i <= r && z[i - l] < r - i + 1) {
                z[i] = z[i - l];
            } else {
                z[i] = Math.max(0, r - i + 1);
                while (i + z[i] < n && s[z[i]] == s[i + z[i]]) ++z[i];
            }
            if (i + z[i] - 1 > r) {
                l = i;
                r = i + z[i] - 1;
            }
        }
        return z;
    }

2223. 构造字符串的总得分和

public class Solution2223 {
    public long sumScores(String s) {
        int n = s.length();
        int[] z = z_function(n, s.toCharArray());

        long cnt = n;
        for (int i : z) {
            cnt += i;
        }
        return cnt;
    }

    private int[] z_function(int n, char[] s) {
        int[] z = new int[n];
        for (int i = 1, l = 0, r = 0; i < n; ++i) {
            if (i <= r && z[i - l] < r - i + 1) {
                z[i] = z[i - l];
            } else {
                z[i] = Math.max(0, r - i + 1);
                while (i + z[i] < n && s[z[i]] == s[i + z[i]]) ++z[i];
            }
            if (i + z[i] - 1 > r) {
                l = i;
                r = i + z[i] - 1;
            }
        }
        return z;
    }
}

3031. 将单词恢复初始状态所需的最短时间 II

public class Solution3031 {
    public int minimumTimeToInitialState(String word, int k) {
        int n = word.length();
        int[] z = z_function(n, word.toCharArray());

        for (int i = 0; i < n; i += k) {
            if (i + k < n && z[i + k] + (i + k) >= n) {
                return (i + k) / k;
            }
        }
        return (n + k - 1) / k;
    }

    private int[] z_function(int n, char[] s) {
        int[] z = new int[n];
        for (int i = 1, l = 0, r = 0; i < n; ++i) {
            if (i <= r && z[i - l] < r - i + 1) {
                z[i] = z[i - l];
            } else {
                z[i] = Math.max(0, r - i + 1);
                while (i + z[i] < n && s[z[i]] == s[i + z[i]]) ++z[i];
            }
            if (i + z[i] - 1 > r) {
                l = i;
                r = i + z[i] - 1;
            }
        }
        return z;
    }
}

3045. 统计前后缀下标对 II

public class Solution3045 {
    public long countPrefixSuffixPairs(String[] words) {
        TrieNode root = new TrieNode();
        long ans = 0;
        for (String word : words) {
            int n = word.length();
            int[] z = z_function(n, word.toCharArray());
            z[0] = n;

            TrieNode node = root;
            for (int i = 0; i < word.length(); i++) {
                int idx = word.charAt(i) - 'a';
                if (node.child[idx] == null) {
                    break;
                }
                node = node.child[idx];
                if (i + 1 == z[n - 1 - i]) {
                    ans += node.cnt;
                }
            }

            // insert
            node = root;
            for (char c : word.toCharArray()) {
                int idx = c - 'a';
                if (node.child[idx] == null) {
                    node.child[idx] = new TrieNode();
                }
                node = node.child[idx];
            }
            node.cnt++;
        }
        return ans;
    }

    static class TrieNode {
        TrieNode[] child;
        int cnt;

        public TrieNode() {
            this.child = new TrieNode[26];
            this.cnt = 0;
        }
    }

    private int[] z_function(int n, char[] s) {
        int[] z = new int[n];
        for (int i = 1, l = 0, r = 0; i < n; ++i) {
            if (i <= r && z[i - l] < r - i + 1) {
                z[i] = z[i - l];
            } else {
                z[i] = Math.max(0, r - i + 1);
                while (i + z[i] < n && s[z[i]] == s[i + z[i]]) ++z[i];
            }
            if (i + z[i] - 1 > r) {
                l = i;
                r = i + z[i] - 1;
            }
        }
        return z;
    }
}

（全文完）