Доктора Кнут, Моррис и Пратт, или Как я перестал бояться и полюбил префикс-функцию

Если вы не знаете, что такое префикс-функция строки, не знаете, как она вычисляется, или, что самое главное, не до конца понимаете, почему алгоритм вычисления префикс-функции работает за линейное время, то эта статья для вас.

Я прошел через череду осознаний и озарений, прежде чем достичь просветления, и теперь предлагаю вам пройти этот путь вместе со мной.

Этап 1. Определение

Первое мое знакомство с префикс-функцией произошло еще в школе. Я готовился к олимпиадам по программированию, и конечно же в моем "джентльменском наборе" подготовки был алгоритм Кнута-Морриса-Пратта, который позволяет найти подстроку длины в строке длины за времени.

Итак, что такое префикс-функция? Прежде чем перейти к ответу на этот вопрос, давайте познакомимся с функцией . Эта функция отображает строку в число, которое является длиной наибольшего суффикса строки , который совпадает с ее префиксом той же длины, но при этом не совпадает со всей строкой .

Без примеров сразу может быть непонятно. Давайте рассмотрим, например, строку . У этой строки есть префиксы , , , , и суффиксы , , , , . Пересечением префиксов и суффиксов являются строки , и . Самая длинная из этих строк, которая не является исходной строкой, - это строка , ее длина равна 3. Соответственно, .

Другие примеры: , , .

Написать функцию, которая вычисляет можно, например, так (язык C++):

int pi(string s) {
    int n = s.size();
    for (int suffix_len = n - 1; suffix_len > 0; suffix_len--) {
        string suffix = s.substr(n - suffix_len, suffix_len);
        string prefix_same_len = s.substr(0, suffix_len);
        if (suffix == prefix_same_len) {
            return suffix_len;
        }
    }
    return 0;
}

Мы проходим по всем суффиксам строки по убыванию их длины, для каждого суффикса находим префикс той же длины и сравниваем их. Как только встретим совпадающие суффикс и префикс - возвращаем их длину. Такая функция, очевидно, будет работать за , но нас это пока что не расстраивает, на этом этапе мы просто хотим разобраться с определением.

Теперь же, когда мы узнали и осознали, что такое , будет гораздо проще разобраться, что такое префикс-функция. Давайте обозначим ее . Так вот, префикс-функция строки - это массив чисел , где - префикс строки до -го символа включительно, .

То есть префикс-функция - это массив чисел для всех префиксов строки (количество префиксов строки совпадает с ее длиной, то есть размер массива равен длине строки).

Например , , , , .

Еще пример: .

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

Этап 2. Применение

Как же префикс-функция может помочь находить подстроку в строке? Идея довольно красивая: давайте склеим подстроку и строку, вставив между ними символ, который гарантированно не будет встречаться ни в подстроке, ни в строке. Далее, посчитаем префикс-функцию для получившейся склеенной строки. Тогда, если в какой-то момент значение префикс-функции станет равно длине искомой подстроки, это будет означать, что мы нашли вхождение подстроки в строку.

Пример: ищем подстроку в строке . Предполагаем, что символ не встречается во входном алфавите. Склеив строки, получаем строку .

Ее префикс-функция равна . Видим, что в функции мы дважды получили , которое является длиной искомой подстроки. Следовательно, в этих позициях заканчиваются вхождения подстроки в исходную строку.

Объясню немного подробнее. Пусть длина искомой подстроки равна . Мы нашли такой префикс склеенной строки , для которого . То есть мы нашли суффикс длины , совпавший с префиксом длины . Префикс длины - это и есть наша искомая подстрока. Суффикс длины гарантированно полностью лежит в исходной строке, в которой мы ищем вхождения. Гарантирует нам это добавленный служебный символ . Благодаря нему, префиксы и суффиксы склеенной строки не пересекаются (за исключением самой строки). Следовательно, найденный суффикс является вхождением подстроки в строку.

Что ж, если мы научимся считать префикс-функцию за линейное время от длины строки, то мы получим алгоритм поиска подстроки длины в строке длины , работающий за времени.

Этап 3. Непонимание

Еще будучи в школе, я узнал, что префикс-функцию можно считать за линейное время. Без предисловий приведу алгоритм, который это делает:

vector<int> prefix_function(string s) {
    int n = s.size();
    vector<int> pi(n);
    pi[0] = 0;
    for (int i = 1; i < n; i++) {
        int k = pi[i - 1];
        while (k > 0 && s[i] != s[k]) {
            k = pi[k - 1];
        }
        if (s[i] == s[k]) {
            k++;
        }
        pi[i] = k;
    }
    return pi;
}

Здесь pi[i] - это в наших обозначениях. То, что pi[0] = 0 - довольно очевидно, потому что функция от строки длины всегда равна .

Сложности в понимании происходящего начинаются, когда мы оказываемся в цикле:

• По какой логике меняется значение k ? И почему в итоге k оказывается равно pi[i] ?

• Почему, несмотря на цикл while внутри цикла for , алгоритм работает за линейное время?

Когда я был юным школьником, я попытался разобраться в этих вопросах, но не преуспел, поскольку не был достаточно заинтересован и замотивирован. Только в вузе я осознал всю прелесть строгих математических доказательств и перестал принимать на веру недоказанные утверждения (если они не аксиомы). А в школе я просто поверил, что эта функция работает, работает быстро (тесты же проходят), зазубрил, как она пишется, и периодически применял при решении задач.

В студенческие времена разобраться с префикс-функцией руки у меня все никак не доходили. И только став преподавателем по программированию, в курсе которого фигурирует алгоритм Кнута-Морриса-Пратта, я нашел в себе достаточно мотивации, чтобы наконец расставить все точки над pi .

Теперь же, если у вас не осталось никаких других вопросов, кроме двух выше написанных, давайте попробуем на них ответить.

Этап 4. Наблюдение

Я разбирался с префикс-функцией, используя две статьи: на e-maxx и на neerc.ifmo. Во многом они пересекаются и одинаковым образом выводят линейный алгоритм. Далее последует просто более подробный пересказ этих статей, где я постараюсь разобрать каждый пункт с примерами и рисунками, которые помогли в понимании лично мне.

В обеих статьях авторы делают следующее наблюдение, которое помогает вывести алгоритм:

.

Или, на языке кода: pi[i + 1] <= pi[i] + 1 .

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

Доказывается это от противного. Предположим, что возможен случай, при котором

pi[i + 1] > pi[i] + 1 .

Или, другими словами, pi[i] < pi[i + 1] - 1 .

Для того, чтобы было выполнено это неравенство, pi[i + 1] должно быть равно как минимум , иначе мы получим значение pi[i] меньше нуля, что невозможно по определению.

Итак, пусть pi[i + 1] = k , k >= 2 .

То есть суффикс длины k совпадает с префиксом длины k . Однако тогда, сделав шаг назад, можно сказать, что у предыдущего префикса суффикс длины k - 1 совпадает с префиксом длины k - 1 . На рисунке ниже: из того, что совпадают красные области, следует то, что совпадают и желтые.

Это означает, что pi[i] >= k - 1 , то есть pi[i] >= pi[i + 1] - 1 . Пришли к противоречию.

Как это наблюдение поможет нам считать префикс-функцию? Предположим, мы находимся на шаге i , уже вычислив pi для всех значений меньше i . Тогда, исходя из наблюдения, у нас возможны два варианта:

• pi[i] = pi[i - 1] + 1 ,

• pi[i] <= pi[i - 1] .

В каком случае будет выполнено первое равенство? Только в том случае, если префикс и суффикс, совпадавшие на предыдущем шаге, увеличились и продолжили совпадать на текущем. То есть, если s[i] == s[pi[i - 1]] . На рисунке ниже: желтая область увеличивается до красной, только если зеленые символы совпадают.

Давайте напишем код, который использует это наблюдение. Если сравниваемые символы не совпали - будем вызывать медленную функцию, написанную на первом этапе. Только теперь мы можем быть уверены, что при вызове медленной функции pi[i] не будет превышать pi[i - 1] . Можем передавать это число в функцию, чтобы ограничить перебор:

int pi_slow(string s, int max_suffix_len) {
    int n = s.size();
    for (int suffix_len = max_suffix_len; suffix_len > 0; suffix_len--) {
        string suffix = s.substr(n - suffix_len, suffix_len);
        string prefix_same_len = s.substr(0, suffix_len);
        if (suffix == prefix_same_len) {
            return suffix_len;
        }
    }
    return 0;
}

vector<int> prefix_function(string s) {
    int n = s.size();
    vector<int> pi(n);
    pi[0] = 0;
    for (int i = 1; i < n; i++) {
        int k = pi[i - 1];
        if (s[i] == s[k]) {
            k++;
        } else {
            k = pi_slow(s.substr(0, i + 1), k);
        }
        pi[i] = k;
    }
    return pi;
}

Я специально обозначил pi[i - 1] за k , чтобы нейминг был таким же, как в коде, к которому мы хотим прийти в конце (см. этап 3). Теперь становится понятнее, откуда в том коде взялось условие

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

Теперь нам осталось разобраться, как ускорить случай, когда

s[i] != s[pi[i - 1]] .

Этап 5. Оптимизация

Итак, мы находимся в такой ситуации (зеленый и синий символы не совпали):

Теперь мы хотим найти такой максимальный префикс, который совпадет с суффиксом, заканчивающимся на s[i] . Мы можем заметить, что такой префикс, за исключением последнего символа, будет полностью совпадать с каким-то суффиксом желтой области. На рисунке ниже: фиолетовые области совпадают, pi[i] будет равно длине фиолетовой области вместе с синим символом.

Теперь, если мы научимся быстро находить максимальную фиолетовую область и проверять, не равен ли следующий после нее символ s[i] (не является ли он синим), то мы существенно оптимизируем наш перебор. Как же нам быстро находить такую фиолетовую область? Попробуйте понять сами, прежде чем читать дальше.

Ключ к пониманию кроется в том, что не нужно забывать о равенстве желтых областей, нужно использовать этот факт. Мы можем дополнить рисунок следующим образом:

На прошлом рисунке было видно, что фиолетовая область есть и в начале, и в конце желтой. Теперь мы просто нарисовали ее в обеих желтых областях с двух сторон. После этого становится очевидным, что максимальная фиолетовая область будет равна функции от желтой области (по определению функции ).

Мы эту функцию уже посчитали на одном из прошлых шагов, потому что желтая область является одним из префиксов нашей строки. Желтая область, являющаяся префиксом, заканчивается на индексе pi[i - 1] - 1 , поэтому функция от нее лежит в

pi[pi[i - 1] - 1] .

Что же нам делать, если синие символы не совпали? Искать следующую подходящую область!

Вспомним, что за k мы обозначали pi[i - 1] , тогда переход к следующей подходящей области будет выглядеть, как k = pi[k - 1] . Оставлю последний рисунок, прежде чем у меня закончатся цвета:

До каких пор мы можем уменьшать k ? Пока оно не станет равным . Когда k станет равно нулю, нужно проверить, не равен ли первый символ строки синему (s[i] ). Если равен, то тогда pi[i] будет равно .

Итого, получаем следующий код:

vector<int> prefix_function(string s) {
    int n = s.size();
    vector<int> pi(n);
    pi[0] = 0;
    for (int i = 1; i < n; i++) {
        int k = pi[i - 1];
        while (k > 0) {
            if (s[i] == s[k]) {
                k++;
                break;
            }
            k = pi[k - 1];
        }
        if (k == 0 && s[i] == s[0]) {
            k = 1;
        }
        pi[i] = k;
    }
    return pi;
}

Если внимательно сравнить его с тем кодом, который я приводил в этапе 3, то можно понять, что они полностью эквиваленты:

vector<int> prefix_function(string s) {
    int n = s.size();
    vector<int> pi(n);
    pi[0] = 0;
    for (int i = 1; i < n; i++) {
        int k = pi[i - 1];
        while (k > 0 && s[i] != s[k]) {
            k = pi[k - 1];
        }
        if (s[i] == s[k]) {
            k++;
        }
        pi[i] = k;
    }
    return pi;
}

В итоговой версии внутри while не происходит увеличение k , оно вынесено в конец и объединено с проверкой в случае k == 0 .

Итак, я надеюсь, у меня получилось ответить на первый из вопросов, сформулированных в этапе 3:

По какой логике меняется значение k ? И почему в итоге k оказывается равно pi[i] ?

Остался последний!

Этап 6. Доказательство линейности

Почему, несмотря на цикл while внутри цикла for , алгоритм работает за линейное время?

На самом деле, ответ на этот вопрос гораздо легче, чем на первый. Однако все доказательства линейности, которые я видел, были сформулированы так, что у меня пухла голова, когда я пытался их осознать.

Сейчас я постараюсь сформулировать ту же мысль, которая есть во многих других источниках, но другими словами, более понятными лично мне. Надеюсь, вам тоже это поможет.

Итак, внутри цикла for мы занимаемся тем, что увеличиваем и уменьшаем переменную k . Причем увеличиваем мы ее всегда только на единицу за один проход, а уменьшаем несколько раз внутри цикла while на неизвестную величину.

На всякий случай поясню, почему внутри while переменная k обязательно уменьшается: pi[k - 1] < k , потому что по определению (функция меньше длины входной строки).

Далее заметим, что в конце тела цикла мы пишем pi[i] = k , а в начале следующей итерации - k = pi[i - 1] . То есть между итерациями значение k не изменяется. Это очень важное наблюдение, отсутствие которого долго мешало мне осознать доказательство.

Из этого наблюдения мы делаем вывод, что переменная k за все итерации увеличится максимум на , потому что всего у нас итераций, в каждой увеличение происходит максимум на , а между итерациями переменная k не изменяется. То есть максимальное возможное значение k - это (изначально у нас k = pi[0] = 0).

Отсюда следует то, что уменьшиться переменная k может максимум те же раз. Потому что в противном случае она оказалась бы меньше нуля, что невозможно.

И отсюда уже следует то, что суммарное количество итераций внутреннего цикла while не превышает , потому что внутри каждой итерации мы гарантированно уменьшаем k . То есть все итерации while в сумме выполнятся за , а значит асимптотика всего алгоритма будет равна .

Что и требовалось доказать.

Этап 7. Заключение

На этом наш путь к просветлению подходит к концу.

Я надеюсь, что все написанное в статье было правдивым и понятным.

Спасибо, что дочитали до конца!

UPDATE

Как мне верно заметили в комментариях, переменную k можно не переобъявлять на каждой итерации внешнего цикла, а объявить один раз перед циклом:

vector<int> prefix_function(string s) {
    int n = s.size();
    vector<int> pi(n);
    pi[0] = 0;
    int k = 0;
    for (int i = 1; i < n; i++) {
        while (k > 0 && s[i] != s[k]) {
            k = pi[k - 1];
        }
        if (s[i] == s[k]) {
            k++;
        }
        pi[i] = k;
    }
    return pi;
}

В таком виде линейность алгоритма становится ещё более наглядной.