Z algorithm (lib/string/z_algorithm.hpp)

View this file on GitHub
Last update: 2025-07-08 15:34:19+09:00
Include: #include "lib/string/z_algorithm.hpp"
Link: https://snuke.hatenablog.com/entry/2014/12/03/214243

Verified with

test/yosupo/string/zalgorithm.test.cpp

Code

#pragma once
#include <vector>

/// @brief Z algorithm
/// @details Z[i] := S[i:Z[i]] == S[0:Z[i]-i]
/// @see https://snuke.hatenablog.com/entry/2014/12/03/214243
template <class Container>
std::vector<int> z_algorithm(const Container &s) {
    int n = s.size();
    std::vector<int> res(n);
    res[0] = n;
    int i = 1, j = 0;
    while (i < n) {
        while (i + j < n && s[j] == s[i + j]) ++j;
        res[i] = j;
        if (!j) {
            ++i;
            continue;
        }
        int k = 1;
        while (i + k < n && k + res[k] < j) res[i + k] = res[k], ++k;
        i += k, j -= k;
    }
    return res;
}

#line 2 "lib/string/z_algorithm.hpp"
#include <vector>

/// @brief Z algorithm
/// @details Z[i] := S[i:Z[i]] == S[0:Z[i]-i]
/// @see https://snuke.hatenablog.com/entry/2014/12/03/214243
template <class Container>
std::vector<int> z_algorithm(const Container &s) {
    int n = s.size();
    std::vector<int> res(n);
    res[0] = n;
    int i = 1, j = 0;
    while (i < n) {
        while (i + j < n && s[j] == s[i + j]) ++j;
        res[i] = j;
        if (!j) {
            ++i;
            continue;
        }
        int k = 1;
        while (i + k < n && k + res[k] < j) res[i + k] = res[k], ++k;
        i += k, j -= k;
    }
    return res;
}