Что такое Z-функция?

Z-функция строки s — это массив z.

z[i] — это длина самого длинного совпадения строки, которая начинается в позиции i, с префиксом всей строки.

Обычно:

z[0] = 0

Простой пример

Строка:

aaaaa

Посмотрим на каждую позицию:

i = 1: aaaa совпадает с prefix aaaa → z[1] = 4
i = 2: aaa совпадает с prefix aaa → z[2] = 3
i = 3: aa совпадает с prefix aa → z[3] = 2
i = 4: a совпадает с prefix a → z[4] = 1

Поэтому:

s:  a a a a a
i:  0 1 2 3 4
z:  0 4 3 2 1

Еще пример

Строка:

abacaba

Z-функция:

s:  a b a c a b a
i:  0 1 2 3 4 5 6
z:  0 0 1 0 3 0 1

Почему z[4] = 3?

s[4...] = aba
prefix = aba

совпало 3 символа

Отличие от префикс-функции

Наивный способ

Для каждой позиции i можно просто сравнивать символы с началом строки.

vector<int> z_function_slow(string s) {
    int n = s.size();

    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[i], можно использовать уже известную информацию, если i попадает внутрь [l, r].

Что делать, если i вне [l, r]?

Если:

i > r

Значит про эту позицию мы ничего не знаем. Тогда просто начинаем сравнивать с нуля:

while (i + z[i] < n && s[z[i]] == s[i + z[i]]) {
    z[i]++;
}

После этого, если получилось хорошее совпадение, обновляем [l, r].

Что делать, если i внутри [l, r]?

Если:

i <= r

Тогда часть ответа уже можно взять из ранее посчитанного значения.

Позиция i внутри блока [l, r]. Ей соответствует позиция:

i - l

Поэтому сначала ставим:

z[i] = min(r - i + 1, z[i - l]);

Основной код Z-функции

vector<int> z_function(string s) {
    int n = s.size();

    vector<int> z(n);

    int l = 0;
    int r = 0;

    for (int i = 1; i < n; i++) {
        if (i <= r) {
            z[i] = min(r - i + 1, z[i - l]);
        }

        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;
}

Полная программа

#include <bits/stdc++.h>
using namespace std;

vector<int> z_function(string s) {
    int n = s.size();

    vector<int> z(n);

    int l = 0;
    int r = 0;

    for (int i = 1; i < n; i++) {
        if (i <= r) {
            z[i] = min(r - i + 1, z[i - l]);
        }

        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;
}

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

    string s;
    cin >> s;

    vector<int> z = z_function(s);

    for (int x : z) {
        cout << x << ' ';
    }

    return 0;
}

Input / Output

Input

abacaba

Output

0 0 1 0 3 0 1

Разбор переменных l и r

[l, r] — это самый правый найденный отрезок, который совпадает с префиксом.

Например, если:

s = abacaba

На позиции i = 4 есть совпадение:

s[4...6] = aba
s[0...2] = aba

Тогда:

l = 4
r = 6

Почему алгоритм O(n)?

В коде есть while, но суммарно он работает не больше O(n).

Каждый раз, когда while делает успешное сравнение и увеличивает z[i], правая граница r двигается вправо.

Поиск подстроки через Z-функцию

Нужно найти все вхождения pattern в text.

Строим строку:

pattern + "#" + text

Потом считаем Z-функцию.

Символ # должен быть таким, которого нет в строках.

Код поиска подстроки

#include <bits/stdc++.h>
using namespace std;

vector<int> z_function(string s) {
    int n = s.size();

    vector<int> z(n);

    int l = 0;
    int r = 0;

    for (int i = 1; i < n; i++) {
        if (i <= r) {
            z[i] = min(r - i + 1, z[i - l]);
        }

        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;
}

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

    string pattern, text;

    cin >> pattern >> text;

    string s = pattern + "#" + text;

    vector<int> z = z_function(s);

    int m = pattern.size();

    for (int i = m + 1; i < (int)s.size(); i++) {
        if (z[i] == m) {
            cout << i - m - 1 << ' ';
        }
    }

    return 0;
}

Пример поиска

Input

aba
abacaba

Output

0 4

Потому что aba встречается в abacaba на позициях:

0 и 4

Почему позиция i - m - 1?

Общая строка:

pattern + "#" + text

Перед text стоит:

m символов pattern + 1 символ #

Поэтому позиция i в общей строке соответствует позиции:

i - m - 1

в строке text.

Функция find_occurrences

Можно сделать отдельную функцию, которая возвращает все позиции pattern в text.

vector<int> find_occurrences(string pattern, string text) {
    string s = pattern + "#" + text;

    vector<int> z = z_function(s);
    vector<int> ans;

    int m = pattern.size();

    for (int i = m + 1; i < (int)s.size(); i++) {
        if (z[i] == m) {
            ans.push_back(i - m - 1);
        }
    }

    return ans;
}

Сравнение с КМП

Поиск всех border строки

Border строки — это строка, которая одновременно является префиксом и суффиксом, но не равна всей строке.

Через Z-функцию можно найти все длины border.

Если:

i + z[i] == n

Значит suffix, который начинается в i, совпал с prefix длины z[i].

Код: все border через Z-функцию

vector<int> get_borders(string s) {
    int n = s.size();

    vector<int> z = z_function(s);
    vector<int> borders;

    for (int i = 1; i < n; i++) {
        if (i + z[i] == n) {
            borders.push_back(z[i]);
        }
    }

    sort(borders.begin(), borders.end());

    return borders;
}

Пример border

Input

ababa

Z-функция:

0 0 3 0 1

Border длины:

1 3

Это строки:

a
aba

Минимальный период строки

Период строки — это длина блока, который повторяется и образует всю строку.

Например:

ababab

Минимальный период:

2

Потому что:

ab ab ab

Период через Z-функцию

Перебираем длину периода p. Если остаток строки с позиции p совпадает с префиксом длины n - p, то p может быть периодом.

Условие:

z[p] >= n - p

И еще нужно:

n % p == 0

Код: минимальный период через Z

int minimal_period(string s) {
    int n = s.size();

    vector<int> z = z_function(s);

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

    return n;
}

Подсчет вхождений каждого префикса

Z-функция показывает, где prefix строки встречается внутри строки.

Если z[i] = len, то в позиции i встречаются все prefix длины:

1, 2, 3, ..., len

Для быстрого подсчета можно использовать разностный массив.

Код: сколько раз встречается каждый prefix

vector<int> count_prefixes(string s) {
    int n = s.size();

    vector<int> z = z_function(s);
    vector<int> diff(n + 2);

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

    vector<int> cnt(n + 1);

    for (int len = 1; len <= n; len++) {
        cnt[len] = cnt[len - 1] + diff[len];
    }

    for (int len = 1; len <= n; len++) {
        cnt[len]++;
    }

    return cnt;
}

Асимптотика

Когда использовать Z-функцию?

Типичные ошибки

• Думать, что z[i] считается для строки s[0...i], как prefix function
• Забывать, что z[i] сравнивает s[i...] с prefix всей строки
• Ошибаться в формуле z[i] = min(r - i + 1, z[i - l])
• Забывать обновить l и r
• Использовать разделитель, который есть в pattern или text
• Ошибаться в позиции i - m - 1
• Путать z[0] = 0 и z[0] = n
• В периоде забыть проверку n % p == 0

Короткий шаблон

vector<int> z_function(string s) {
    int n = s.size();

    vector<int> z(n);

    int l = 0;
    int r = 0;

    for (int i = 1; i < n; i++) {
        if (i <= r) {
            z[i] = min(r - i + 1, z[i - l]);
        }

        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;
}

Короткий шаблон поиска pattern

vector<int> find_occurrences(string pattern, string text) {
    string s = pattern + "#" + text;

    vector<int> z = z_function(s);
    vector<int> ans;

    int m = pattern.size();

    for (int i = m + 1; i < (int)s.size(); i++) {
        if (z[i] == m) {
            ans.push_back(i - m - 1);
        }
    }

    return ans;
}

Главное запомнить