基于Z函数的字符串前缀出现次数统计算法超时优化咨询

问题描述 #

给定字符串s，对于每个i（1≤i≤|s|），求长度为i的前缀在s中的出现次数。

输入格式 #

• 第一行输入整数q（1≤q≤10^5）——测试用例数。
• 每个测试用例为字符串s，长度1到10^{6，仅含小写拉丁字母。所有用例字符串总长度不超过10}6。

输出格式 #

每个测试用例输出|s|个整数c₁,c₂,…,c_|s|，其中c[i]是长度为i的前缀在s中的出现次数。

示例

输入：

5
abacaba
eeeee
abcdef
ababababa
kekkekkek

输出：

4 2 2 1 1 1 1 
5 4 3 2 1 
1 1 1 1 1 1 
5 4 4 3 3 2 2 1 1 
6 3 3 2 2 2 1 1 1

要求必须使用Z函数解决，总处理时间不超过2秒。

初始解法 #

我最初的实现代码如下：

#include <iostream>
#include <string>
#include <vector>

std::vector<int> ZFunc(const std::string& s) {
    const int sz = s.size();
    std::vector<int> z(sz, 0);

    for (int i = 1, l = 0, r = 0; i != sz; ++i) {
        if (r >= i)
            z[i] = std::min(z[i - l], r - i + 1);

        while (z[i] + i < sz && s[i + z[i]] == s[z[i]])
            z[i]++;

        if (z[i] > r - i + 1) {
            l = i;
            r = i + z[i] - 1;
        }
    }

    return z;
}

int main() {
    int n;
    std::cin >> n;

    std::vector<std::vector<int>> res(n);

    for (int k = 0; k != n; ++k) {
        std::string s;
        std::cin >> s;

        res[k].resize(s.size(), 1);
        std::vector<int> z = ZFunc(s);

        for (int i = 1; i != z.size(); ++i) {
            while (z[i]--)
                res[k][z[i]]++;
        }
    }

    for (const auto& ivec : res) {
        for (int i : ivec)
            std::cout << i << " ";
        std::cout << std::endl;
    }

    return 0;
}

问题分析 #

这个解法存在超时问题，核心原因是这段代码：

while (z[i]--)
    res[k][z[i]]++;

当遇到类似aaaaa...aaabcd这类字符串时，Z函数的z[i]值会很大（比如前n-1个位置的z[i]都是n-i），此时这个while循环的总执行次数会达到O(n²)，远超时间限制。

优化方案 #

要解决这个问题，我们可以用**差分数组（前缀和技巧）**来替代逐个累加的操作，把区间增量操作的复杂度从O(z[i])降到O(1)。

具体思路：

1. 初始化一个差分数组diff，大小为s.size() + 2（避免处理边界时越界），初始值全为0。
2. 遍历Z数组的每个值z_val = z[i]（i从1开始）：
   • 如果z_val > 0，说明长度1到z_val的前缀各多出现一次，我们只需要在差分数组中标记：diff[1] += 1，diff[z_val + 1] -= 1。
3. 计算差分数组的前缀和，得到每个长度k对应的增量次数。
4. 每个长度k的最终出现次数就是初始的1加上对应的增量次数。

修改后的代码如下：

#include <iostream>
#include <string>
#include <vector>

std::vector<int> ZFunc(const std::string& s) {
    const int sz = s.size();
    std::vector<int> z(sz, 0);

    for (int i = 1, l = 0, r = 0; i != sz; ++i) {
        if (r >= i)
            z[i] = std::min(z[i - l], r - i + 1);

        while (z[i] + i < sz && s[i + z[i]] == s[z[i]])
            z[i]++;

        if (z[i] > r - i + 1) {
            l = i;
            r = i + z[i] - 1;
        }
    }

    return z;
}

int main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);

    int q;
    std::cin >> q;

    while (q--) {
        std::string s;
        std::cin >> s;
        int n = s.size();
        std::vector<int> z = ZFunc(s);
        
        std::vector<int> diff(n + 2, 0); // 索引从1到n
        for (int i = 1; i < n; ++i) {
            int z_val = z[i];
            if (z_val > 0) {
                diff[1] += 1;
                if (z_val + 1 <= n) {
                    diff[z_val + 1] -= 1;
                }
            }
        }

        // 计算前缀和得到增量，同时输出结果
        int current = 0;
        for (int k = 1; k <= n; ++k) {
            current += diff[k];
            std::cout << (1 + current) << " ";
        }
        std::cout << "\n";
    }

    return 0;
}

额外优化：加入std::ios::sync_with_stdio(false);和std::cin.tie(nullptr);来加速输入输出，避免因大量IO操作导致的超时。

这样修改后，整个算法的时间复杂度就变成了O(total_length)，完全符合题目要求的时间限制。

内容的提问来源于stack exchange，提问作者vadyaov