Префикс-функция

Что такое префикс?

Префикс строки — это начало строки.

Например, строка:

abcde

Ее префиксы:

a
ab
abc
abcd
abcde

Пустую строку тоже иногда считают префиксом, но в префикс-функции обычно смотрят непустые префиксы.

Что такое суффикс?

Суффикс строки — это конец строки.

Например, строка:

abcde

Ее суффиксы:

e
de
cde
bcde
abcde

Что такое префикс-функция?

Префикс-функция для строки s — это массив pi.

pi[i] — это длина наибольшего собственного префикса строки s[0...i], который одновременно является суффиксом этой же строки.

берем подстроку s[0...i]
ищем самый длинный префикс
который равен суффиксу
но не равен всей строке

Собственный префикс — это префикс, который не равен всей строке.

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

Строка:

ababa

Рассмотрим всю строку ababa.

Ее префиксы:

a
ab
aba
abab
ababa

Ее суффиксы:

a
ba
aba
baba
ababa

Самый длинный собственный префикс, который равен суффиксу:

aba

Значит для последней позиции:

pi[4] = 3

Пример массива pi

Для строки:

ababa

Префикс-функция:

s:  a b a b a
i:  0 1 2 3 4
pi: 0 0 1 2 3

Объяснение:

pi[0] = 0, потому что у одной буквы нет собственного непустого префикса
pi[1] = 0, строка ab, нет совпадения
pi[2] = 1, строка aba, совпадает a
pi[3] = 2, строка abab, совпадает ab
pi[4] = 3, строка ababa, совпадает aba

Еще пример

Строка:

aaaaa

Префикс-функция:

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

Почему так? Потому что каждая следующая буква продолжает совпадение.

Идея алгоритма

Хотим посчитать pi[i]. Пусть для предыдущей позиции уже известно значение:

j = pi[i - 1]

Это значит, что у строки s[0...i-1] уже есть совпавший префикс длины j.

Теперь хотим добавить символ s[i].

если s[i] == s[j], то совпадение увеличивается
если s[i] != s[j], надо откатиться на меньший префикс

Что значит откат?

Если символы не совпали, мы не начинаем проверку с нуля. Мы используем уже посчитанные значения pi.

Было:

j = pi[i - 1]

Если s[i] != s[j], то делаем:

j = pi[j - 1]

Это значит: переходим к следующей меньшей границе.

Именно из-за таких откатов алгоритм работает быстро.

Основной код префикс-функции

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

    vector<int> pi(n);

    for (int i = 1; i < n; i++) {
        int j = pi[i - 1];

        while (j > 0 && s[i] != s[j]) {
            j = pi[j - 1];
        }

        if (s[i] == s[j]) {
            j++;
        }

        pi[i] = j;
    }

    return pi;
}

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

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

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

    vector<int> pi(n);

    for (int i = 1; i < n; i++) {
        int j = pi[i - 1];

        while (j > 0 && s[i] != s[j]) {
            j = pi[j - 1];
        }

        if (s[i] == s[j]) {
            j++;
        }

        pi[i] = j;
    }

    return pi;
}

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

    string s;
    cin >> s;

    vector<int> pi = prefix_function(s);

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

    return 0;
}

Input / Output

Input

ababa

Output

0 0 1 2 3

Разбор работы алгоритма

Пусть:

s = ababa

Идем по строке:

i = 1, s[i] = b
j = pi[0] = 0
s[1] != s[0]
pi[1] = 0

i = 2, s[i] = a
j = pi[1] = 0
s[2] == s[0]
pi[2] = 1

i = 3, s[i] = b
j = pi[2] = 1
s[3] == s[1]
pi[3] = 2

i = 4, s[i] = a
j = pi[3] = 2
s[4] == s[2]
pi[4] = 3

Сложный пример с откатом

Строка:

abacaba

Префикс-функция:

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

На позиции i = 3 символ c не продолжает префикс aba, поэтому алгоритм делает откат.

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

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

Переменная j может увеличиваться только на 1 за шаг, а при откате уменьшается по значениям pi.

за весь алгоритм j не может увеличиться больше n раз
и не может уменьшиться больше n раз
поэтому суммарно O(n)

Зачем нужна префикс-функция?

Префикс-функция используется в строковых алгоритмах.

1. поиск подстроки в строке
2. алгоритм KMP
3. поиск всех вхождений pattern в text
4. поиск периода строки
5. работа с границами строки

Поиск подстроки через KMP

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

Создаем новую строку:

pattern + "#" + text

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

Если где-то значение pi равно длине pattern, значит pattern закончился в этой позиции.

Код KMP

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

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

    vector<int> pi(n);

    for (int i = 1; i < n; i++) {
        int j = pi[i - 1];

        while (j > 0 && s[i] != s[j]) {
            j = pi[j - 1];
        }

        if (s[i] == s[j]) {
            j++;
        }

        pi[i] = j;
    }

    return pi;
}

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

    string pattern, text;

    cin >> pattern >> text;

    string s = pattern + "#" + text;

    vector<int> pi = prefix_function(s);

    int m = pattern.size();

    for (int i = 0; i < (int)s.size(); i++) {
        if (pi[i] == m) {
            int pos = i - 2 * m;

            cout << pos << ' ';
        }
    }

    return 0;
}

Здесь позиции выводятся с 0-indexed.

Пример KMP

Input

aba
abacaba

Output

0 4

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

0 и 4

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

Мы строим строку:

pattern + "#" + text

Длина pattern равна m. Перед text стоит:

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

Если на позиции i в общей строке найдено совпадение длины m, то начало в text:

i - 2 * m

Граница строки

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

Например:

s = ababa

Границы:

a
aba

Самая длинная граница имеет длину:

pi[n - 1]

Найти все границы строки

Можно начать с:

k = pi[n - 1]

Потом переходить:

k = pi[k - 1]

Так мы найдем все длины границ.

Код: все границы строки

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

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

    vector<int> pi(n);

    for (int i = 1; i < n; i++) {
        int j = pi[i - 1];

        while (j > 0 && s[i] != s[j]) {
            j = pi[j - 1];
        }

        if (s[i] == s[j]) {
            j++;
        }

        pi[i] = j;
    }

    return pi;
}

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

    string s;
    cin >> s;

    vector<int> pi = prefix_function(s);

    vector<int> borders;

    int k = pi[(int)s.size() - 1];

    while (k > 0) {
        borders.push_back(k);
        k = pi[k - 1];
    }

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

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

    return 0;
}

Пример границ

Input

ababa

Output

1 3

Это длины границ:

a
aba

Период строки

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

Например:

ababab

Имеет период:

Потому что:

ab ab ab

Минимальный период через prefix function

Пусть:

n = s.size()
k = pi[n - 1]

Тогда кандидат на минимальный период:

period = n - k

Если n % period == 0, значит это настоящий период. Иначе минимальный период равен n.

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

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

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

    vector<int> pi(n);

    for (int i = 1; i < n; i++) {
        int j = pi[i - 1];

        while (j > 0 && s[i] != s[j]) {
            j = pi[j - 1];
        }

        if (s[i] == s[j]) {
            j++;
        }

        pi[i] = j;
    }

    return pi;
}

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

    string s;
    cin >> s;

    int n = s.size();

    vector<int> pi = prefix_function(s);

    int period = n - pi[n - 1];

    if (n % period == 0) {
        cout << period;
    } else {
        cout << n;
    }

    return 0;
}

Пример периода

Input

ababab

Output

Строка состоит из повторов:

ab ab ab

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

Иногда нужно узнать, сколько раз каждый префикс строки встречается внутри строки.

Сначала считаем prefix function. Потом каждое значение pi[i] увеличивает счетчик соответствующего префикса.

Код: количество вхождений префиксов

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

    vector<int> pi = prefix_function(s);
    vector<int> cnt(n + 1);

    for (int i = 0; i < n; i++) {
        cnt[pi[i]]++;
    }

    for (int i = n; i > 0; i--) {
        cnt[pi[i - 1]] += cnt[i];
    }

    for (int i = 0; i <= n; i++) {
        cnt[i]++;
    }

    return cnt;
}

cnt[len] показывает, сколько раз префикс длины len встречается в строке.

Асимптотика

Задача	Время	Память
Префикс-функция	O(n)	O(n)
Поиск pattern в text	O(n + m)	O(n + m)
Все границы строки	O(n)	O(n)
Минимальный период	O(n)	O(n)

Когда использовать prefix function?

Задача	Что использовать
Найти pattern в text	KMP
Найти все границы строки	pi[n - 1], потом откаты
Найти минимальный период	n - pi[n - 1]
Работа с совпадениями префикса и суффикса	prefix function

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

Путать префикс и суффикс
Думать, что pi[i] считается для всей строки, а не для s[0...i]
Забыть, что pi[0] всегда 0
В while писать if и делать только один откат
Забыть условие j > 0 в while
В KMP выбрать разделитель, который встречается в строках
Ошибиться с индексом позиции в KMP
Путать 0-indexed и 1-indexed позиции

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

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

    vector<int> pi(n);

    for (int i = 1; i < n; i++) {
        int j = pi[i - 1];

        while (j > 0 && s[i] != s[j]) {
            j = pi[j - 1];
        }

        if (s[i] == s[j]) {
            j++;
        }

        pi[i] = j;
    }

    return pi;
}

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

pi[i] = длина самого длинного префикса
который равен суффиксу строки s[0...i]
если символы совпали → j++
если не совпали → j = pi[j - 1]
prefix function работает за O(n)
KMP ищет pattern в text за O(n + m)