Алгоритм КМП

Что делает КМП?

КМП — это алгоритм Кнута-Морриса-Пратта. Он ищет все вхождения строки pattern в строку text.

pattern — что ищем
text — где ищем
answer — позиции, где pattern встречается в text

Например:

pattern = aba
text = abacaba

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

0 4

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

Самый простой способ — попробовать поставить pattern в каждую позицию text.

for (int i = 0; i + m <= n; i++) {
    bool ok = true;

    for (int j = 0; j < m; j++) {
        if (text[i + j] != pattern[j]) {
            ok = false;
            break;
        }
    }

    if (ok) {
        cout << i << ' ';
    }
}

Но такой метод работает за:

O(n * m)

Если строки большие, это слишком медленно.

Идея КМП

КМП не начинает сравнение заново после ошибки. Он использует информацию о том, какая часть pattern уже совпала.

Для этого используется:

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

если совпадение сломалось
мы не возвращаемся полностью назад
а откатываемся к меньшему совпавшему префиксу

Префикс и суффикс

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

abcde

префиксы:
a
ab
abc
abcd
abcde

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

abcde

суффиксы:
e
de
cde
bcde
abcde

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

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

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

Пример:

s = ababa

pi = 0 0 1 2 3

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

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

Как КМП использует prefix function?

Чтобы найти pattern в text, создаем новую строку:

pattern + "#" + text

Символ # нужен как разделитель. Он не должен встречаться ни в pattern, ни в text.

Потом считаем prefix function для этой новой строки.

если pi[i] == pattern.size()
значит pattern полностью совпал
и одно вхождение найдено

Пример строки для КМП

Пусть:

pattern = aba
text = abacaba

Тогда строим:

s = aba#abacaba

Дальше считаем prefix function для всей строки s.

Если где-то pi[i] = 3, значит строка aba закончилась в этой позиции.

Формула позиции

Пусть:

m = pattern.size()

Если в общей строке s = pattern + "#" + text на позиции i получилось:

pi[i] = m

Тогда позиция начала вхождения в text:

pos = i - 2 * m

В этом коде позиции считаются с 0.

Почему pos = i - 2 * m?

Общая строка имеет вид:

pattern + "#" + text

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

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

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

i - m + 1

Чтобы перейти к позиции внутри text, вычитаем m + 1:

(i - m + 1) - (m + 1) = i - 2 * m

Код КМП

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

Input / Output

Input

aba
abacaba

Output

0 4

Разбор примера

Ищем:

aba

В строке:

abacaba

Проверяем вручную:

позиция 0: aba совпало
позиция 1: bac не совпало
позиция 2: aca не совпало
позиция 3: cab не совпало
позиция 4: aba совпало

Ответ:

0 4

Функция find_occurrences

Удобно сделать отдельную функцию, которая возвращает все позиции.

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

    vector<int> pi = prefix_function(s);
    vector<int> ans;

    int m = pattern.size();

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

    return ans;
}

Полная программа с функцией

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

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

    vector<int> pi = prefix_function(s);
    vector<int> ans;

    int m = pattern.size();

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

    return ans;
}

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

    string pattern, text;

    cin >> pattern >> text;

    vector<int> ans = find_occurrences(pattern, text);

    for (int pos : ans) {
        cout << pos << ' ';
    }

    return 0;
}

Поиск только первого вхождения

Иногда нужно найти только первую позицию. Если вхождения нет, возвращаем -1.

int find_first(string pattern, string 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) {
            return i - 2 * m;
        }
    }

    return -1;
}

Проверка, есть ли pattern в text

Если нужно только понять, есть ли подстрока:

bool contains(string pattern, string text) {
    return find_first(pattern, text) != -1;
}

Это тоже работает за:

O(n + m)

КМП без создания большой строки

Можно искать pattern в text без строки pattern + "#" + text. Сначала считаем prefix function только для pattern.

Потом идем по text и поддерживаем длину текущего совпадения j.

j — сколько символов pattern сейчас совпало
если text[i] == pattern[j], увеличиваем j
если j == m, найдено вхождение

Код КМП без склейки строк

vector<int> kmp_search(string pattern, string text) {
    vector<int> pi = prefix_function(pattern);
    vector<int> ans;

    int n = text.size();
    int m = pattern.size();

    int j = 0;

    for (int i = 0; i < n; i++) {
        while (j > 0 && text[i] != pattern[j]) {
            j = pi[j - 1];
        }

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

        if (j == m) {
            ans.push_back(i - m + 1);
            j = pi[j - 1];
        }
    }

    return ans;
}

Этот вариант удобен, если не хочется создавать дополнительную строку.

Полный код КМП без склейки

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

vector<int> kmp_search(string pattern, string text) {
    vector<int> pi = prefix_function(pattern);
    vector<int> ans;

    int n = text.size();
    int m = pattern.size();

    int j = 0;

    for (int i = 0; i < n; i++) {
        while (j > 0 && text[i] != pattern[j]) {
            j = pi[j - 1];
        }

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

        if (j == m) {
            ans.push_back(i - m + 1);
            j = pi[j - 1];
        }
    }

    return ans;
}

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

    string pattern, text;

    cin >> pattern >> text;

    vector<int> ans = kmp_search(pattern, text);

    for (int pos : ans) {
        cout << pos << ' ';
    }

    return 0;
}

Почему после найденного вхождения делаем откат?

После того как j == m, pattern полностью найден. Но следующее вхождение может пересекаться с текущим.

Поэтому нельзя просто сделать j = 0. Нужно откатиться так:

j = pi[j - 1];

Это позволяет найти пересекающиеся вхождения.

Пример пересекающихся вхождений

Пусть:

pattern = aaa
text = aaaaa

Вхождения:

0 1 2

Они пересекаются:

aaa..
.aaa.
..aaa

Именно поэтому после найденного pattern делаем откат через pi.

КМП для подсчета количества вхождений

Если нужны не позиции, а только количество:

int count_occurrences(string pattern, string text) {
    vector<int> ans = kmp_search(pattern, text);

    return ans.size();
}

Или можно увеличивать счетчик сразу внутри цикла.

Код только для количества

int count_occurrences(string pattern, string text) {
    vector<int> pi = prefix_function(pattern);

    int n = text.size();
    int m = pattern.size();

    int j = 0;
    int cnt = 0;

    for (int i = 0; i < n; i++) {
        while (j > 0 && text[i] != pattern[j]) {
            j = pi[j - 1];
        }

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

        if (j == m) {
            cnt++;
            j = pi[j - 1];
        }
    }

    return cnt;
}

Сравнение с наивным поиском

Метод	Время	Идея
Наивный поиск	O(n * m)	проверяем каждую позицию заново
КМП	O(n + m)	используем prefix function и откаты

Когда использовать КМП?

нужно найти pattern в text
нужно найти все позиции pattern
нужно посчитать количество вхождений
строки большие
нужен гарантированный O(n + m)

Асимптотика

Часть	Время	Память
prefix_function(pattern)	O(m)	O(m)
поиск в text	O(n)	O(1) кроме ответа
весь КМП	O(n + m)	O(m)
вариант со склейкой	O(n + m)	O(n + m)

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

Забыть, что позиции обычно 0-indexed
Взять разделитель, который встречается в pattern или text
Ошибиться в формуле i - 2 * m
После найденного вхождения сделать j = 0 и потерять пересечения
В prefix function заменить while на if
Забыть условие j > 0 в while
Обращаться к pattern[j], когда pattern пустой
Путать n длину text и m длину pattern

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

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

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

vector<int> kmp_search(string pattern, string text) {
    vector<int> pi = prefix_function(pattern);
    vector<int> ans;

    int n = text.size();
    int m = pattern.size();

    int j = 0;

    for (int i = 0; i < n; i++) {
        while (j > 0 && text[i] != pattern[j]) {
            j = pi[j - 1];
        }

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

        if (j == m) {
            ans.push_back(i - m + 1);
            j = pi[j - 1];
        }
    }

    return ans;
}

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

КМП ищет pattern в text
работает за O(n + m)
использует prefix function
j показывает длину текущего совпадения
при ошибке делаем j = pi[j - 1]
после найденного вхождения тоже делаем откат