Базовые строковые алгоритмы

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

Рассмотрим задачу, которая возникает каждый раз, когда вы делаете ctrl+f:

Есть большой текст $t$. Нужно найти все вхождения строки $s$ в него.

Наивное решение со сравнением всех подстрок $t$ длины $|s|$ со строкой $s$ работает за $O(|t| \cdot |s|)$. Если текст большой, то длинные слова в нем искать становится очень долго.

Для решения этой задачи за линейное время придумали префикс-функцию.

Определение. Префикс-функцией от строки $s$ называется массив $p$, где $p_i$ равно длине самого большого префикса строки $s_0 s_1 s_2 \ldots s_i$, который также является и суффиксом этой строки (не считая всю строку).

Например, самый большой префикс, который равен суффиксу для строки «aataataa» — это «aataa»; префикс-функция для этой строки равна $[0, 1, 0, 1, 2, 3, 4, 5]$.

In [7]:
def slow_prefix_function(s):
    n = len(s)
    p = [0]*n
    for i in range(n):
        prefix = s[:i]
        for l in range(1, i):
            if prefix[:l] == prefix[-l:]:
                p[i] = l
    return p

slow_prefix_function('aataataa')
Out[7]:
[0, 0, 1, 0, 1, 2, 3, 4]

(Этот алгоритм работает за $O(n^3)$, но это только пока.)

Как это поможет решить исходную задачу?

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

Соединим подстроки $s$ и $t$ каким-нибудь символом, который не встречается ни там, ни там. Посмотрим на префикс-функцию получившейся строки $s\#t$.

In [19]:
s = "let it go"
t = """let it go, let it go
can't hold it back anymore
let it go, let it go
turn away and slam the door!"""

print((s + '#' + t).replace('\n', ' '))
print(''.join([str(x) for x in slow_prefix_function(s + '#' + t)]))
let it go#let it go, let it go can't hold it back anymore let it go, let it go turn away and slam the door!
00000000000123456789001234567890000000001000000000000000000123456789001234567890000000000000000100000000000

Видно, что все места, где значения равны 9 (длине S) — это концы вхождений $s$ в текст $t$.

Такой алгоритм (посчитать префикс-функцию от $s\#t$ и посмотреть, в каких позициях она равна $|s|$) называется алгоритмом Кнута-Морриса-Пратта.

Как её быстро считать

Рассмотрим ещё несколько примеров префикс-функций:

In [10]:
for s in ['aaaaa', 'abcdef', 'abacabadava', 'antananarivuantananarivu']:
    print(slow_prefix_function(s))
[0, 0, 1, 2, 3]
[0, 0, 0, 0, 0, 0]
[0, 0, 0, 1, 0, 1, 2, 3, 0, 1, 0]
[0, 0, 0, 0, 1, 2, 1, 2, 1, 0, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]

Можно заметить несколько особенностей:

  • $p_0 = 0$ для всех строк: префикс подстроки из первого символа, который не совпадает со всей строкой — это только пустой префикс;
  • $p_{i+1}$ максимум на единицу превосходит $p_i$: если есть префикс, равный суффиксу строки $s_0 s_1 \ldots s_{i+1}$ длины $p_{i+1}$, то, отбросив последний символ, можно получить и правильный суффикс для строки $s_0 s_1 \ldots s_i$, длина которого будет ровно на единицу меньше.

Хочется придумать алгоритм, как считать префикс-функцию за $O(N)$. Логично это делать с помощью динамики: найти формулу для $p_i$ через предыдущие значения.

Заметим, что $p_{i+1} = p_i + 1$ в том и только том случае, когда $s_{p_i} =s_{i+1}$. Например, в строке $\underbrace{aabaa}t\overbrace{aabaa}$ выделен максимальный префикс, равный суффиксу: $p_{10} = 5$. Если следующий символ равен будет равен $t$, то $p_{11} = p_{10} + 1 = 6$.

Но что происходит, когда $s_{p_i}\neq s_{i+1}$? Пусть следующий символ в этой же строке равен $b$.

  • $\implies$ Тогда длина префикса, равного суффиксу новой строки, будет точно меньше 5.
  • $\implies$ Значит, помимо того, что он является суффиксом «aabaab», префикс является префиксом и подстроки «aabaa».
  • $\implies$ Значит, следующий кандидат на проверку — это значение префикс-функции от «aabaa», то есть $p_4 = 2$.
  • $\implies$ Если $s_2 = s_{11}$ (т. е. новый символ совпадает с идущим после префикса-кандидата), то $p_{11} = p_2 + 1 = 2 + 1 = 3$.

В данном случае это действительно так (нужный префикс — «aab»). Но что делать, если, в общем случае, $p_{i+1} \neq p_{p_i+1} $? Тогда мы проводим такое же рассуждение и получаем нового кандидата, меньшей длины — $p_{p_{p_i}}$. Если и этот не подошел — аналогично проверяем меньшего, пока этот индекс не станет нулевым.

In [20]:
def fast_prefix_function(s):
    n = len(s)
    p = [0]*n
    for i in range(1, n):
        cur = p[i - 1]
        # перебираем префикс-функцию, пока не найдем равный символ
        while s[i] != s[cur] and cur > 0:
            cur = p[cur - 1]
         # если нашли, то значение на единицу больше
        if s[i] == s[cur]:
            p[i] = cur + 1
    return p

fast_prefix_function('abacabadabacabax')
Out[20]:
[0, 0, 1, 0, 1, 2, 3, 0, 1, 2, 3, 4, 5, 6, 7, 0]

Асимптотика

Почему это работает за $O(N)$? В худшем случае этот while может работать $O(N)$ раз за одну итерацию. Но оказывается, что в среднем каждый while работает за $O(1)$ — это называется амортизированной асимптотикой.

Доказательство:

  • Как мы ранее заметили, префикс функция возрастает максимум на единицу.
  • $\implies$ Вырасти она может максимум $n-1$ раз.
  • Если мы зашли в while, то каждая его итерация понижает значение префикс-функции хотя бы на один.
  • $\implies$ Количество понижений не больше количества повышений, то есть $O(n)$.
  • $\implies$ Суммарно итераций цикла while — $O(n)$.

Z-функция

Альтернатива префикс-функции — z-функция (примечание: не «зи», а «зет»). Она немного проще для понимания.

Z-функция от строки $s$ — это такой массив $z$, что $z_i$ равно длине максимальной подстроки, начинающейся с $i$-й позиции, которая равна префиксу $s$.

$$\underbrace{aba}c\overbrace{aba}daba \hspace{1em} (z_4 = 3)$$
In [5]:
def slow_z_function(s):
    n = len(s)
    z = [0]*n
    for i in range(1, n):
        suffix = s[i:]
        for l in range(1, i):
            if s[:l] == suffix[:l]:
                z[i] = l
    return z

for s in ['abacabadaba', 'aabcaabaabca', 'antananarivuantananarivu']:
    print(slow_z_function(s))
[0, 0, 1, 0, 3, 0, 1, 0, 3, 0, 1]
[0, 0, 0, 0, 3, 1, 0, 5, 1, 0, 0, 1]
[0, 0, 0, 2, 0, 2, 0, 1, 0, 0, 0, 0, 11, 0, 0, 2, 0, 2, 0, 1, 0, 0, 0, 0]

Z-функцию можно использовать вместо префикс-функции в алгоритме Кнута-Морриса-Пратта — только теперь нужные позиции будут начинаться c $s$, а не заканчиваться. Осталось научиться её искать за $O(n)$.

Как её быстро считать

Заметим, что:

  • $z_0 = 0$ — из-за договоренности (потому что информации не несет);
  • $z_i = 0 \iff s_i \neq s_0$;
  • $z_i > 0 \iff s_i = s_0$.

Будем идти слева направо и хранить z-блок — самую правую подстроку, равную префиксу, которую мы успели обнаружить. Будем обозначать его границы как $l$ и $r$.

Пусть мы сейчас хотим найти $z_i$, а все предыдущие уже нашли. Если новый, $i$-й символ не лежит левее z-блока, значит он лежит либо в z-блоке, либо правее.

  • Если правее, то мы просто наивно перебором найдем $z_i$ (максимальный отрезок, начинающийся с $s_i$ и равный префиксу), и объявим его новым z-блоком.
  • Если $i$-й элемент лежит внутри z-блока, то мы можем посмотреть на значение $z_{i-l}$ и использовать его, чтобы инициализировать $z_i$ чем-то, возможно, отличным от нуля. Если $z_{i-l}$ «не хватает» до границы $z$-блока, то $z_i = z_{i-l}$. Если он упирается, в границу, то «обрежем» его до известной границы и будем увеличивать на единичку.
In [1]:
def fast_z_function(s):
    n = len(s)
    z = [0]*n
    l = 0
    r = 0
    for i in range(1, n):
        if i <= r:
            z[i] = min(r - i + 1, z[i-l])
        while i + z[i] < n and s[z[i]] == s[i+z[i]]:
            z[i] += 1
        if i + z[i] - 1 > r:
            l = i
            r = i +  z[i] - 1
    return z

Асимптотика. В алгоритме мы делаем столько действий, сколько раз сдвигается правая граница z-блока — а это $O(n)$.

Зачем тогда люди используют префикс-функцию

hz

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

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