Подготовка к алгоритмическим задачам

Z-функция: что это, как считать за O(n) и где применять

10 мин чтения

алгоритмы

строки

python

Коротко

Параметр	Значение
Построение	O(n)
Память	O(n)
Поиск подстроки	O(n + m)

Z-функция — одна из двух базовых строковых конструкций (вторая — префикс-функция из алгоритма КМП). Проще для понимания и часто удобнее в задачах.

Определение

Для строки s длины n массив z определяется так:

z[i] — длина наибольшего общего префикса строки s и её суффикса, начинающегося с позиции i.

Проще говоря: насколько далеко строка, начинающаяся с позиции i, совпадает с началом всей строки.

s  =  a  a  b  a  a  b  a  a
i  =  0  1  2  3  4  5  6  7
z  =  -  1  0  5  1  0  2  1

Разберём несколько значений:

z[1] = 1 — суффикс abaabaa совпадает с aabaabaa только в первом символе a;
z[3] = 5 — суффикс aabaa полностью совпадает с началом aabaa, это 5 символов;
z[2] = 0 — суффикс начинается с b, а строка с a, совпадений нет.

Значение z[0] обычно не определяют (оно равно длине всей строки) и ставят ноль или прочерк.

Наивный подсчёт

Для каждой позиции сравниваем символы, пока совпадают:

def z_function_naive(s: str) -> list[int]:
    n = len(s)
    z = [0] * n

    for i in range(1, n):
        while i + z[i] < n and s[z[i]] == s[i + z[i]]:
            z[i] += 1

    return z

Работает, но в худшем случае O(n²) — например, на строке из одинаковых символов aaaa...a.

Линейный алгоритм: идея Z-блока

Ускорение основано на переиспользовании уже посчитанного.

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

s:   a a b a a b a a
     ├───────┤          ← это префикс, значит s[l..r] == s[0..r-l]
     l       r

Когда мы приходим в позицию i внутри блока, мы уже знаем, что происходит в соответствующей позиции у начала строки. Пусть k = i - l — позиция-двойник в префиксе. Тогда значение z[k] подсказывает ответ:

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

Ограничение r - i + 1 необходимо: за пределами блока информации нет, и там придётся сравнивать честно.

После этого дотягиваем z[i] наивным сравнением и, если блок расширился, обновляем l и r.

Реализация

def z_function(s: str) -> list[int]:
    n = len(s)
    z = [0] * n
    l = r = 0            # границы текущего Z-блока

    for i in range(1, n):
        if i < r:
            # используем ранее посчитанное значение
            z[i] = min(r - i, z[i - l])

        # дотягиваем наивно
        while i + z[i] < n and s[z[i]] == s[i + z[i]]:
            z[i] += 1

        # обновляем блок, если ушли правее
        if i + z[i] > r:
            l, r = i, i + z[i]

    return z


print(z_function('aabaabaa'))
# [0, 1, 0, 5, 1, 0, 2, 1]

Почему это O(n)

Тот же аргумент, что и в других алгоритмах со скользящим окном: правая граница r только растёт и никогда не уменьшается.

Каждая итерация внутреннего while увеличивает r хотя бы на единицу. Значит суммарное число таких итераций за всю работу алгоритма не превышает n. Внешний цикл тоже n. Итого линейно.

Это стандартный вопрос на собеседовании — стоит уметь проговорить.

Применение: поиск подстроки

Главное практическое применение. Чтобы найти все вхождения pattern в text, склеим их через разделитель, которого нет ни в одной строке:

def find_all(text: str, pattern: str) -> list[int]:
    if not pattern:
        return []

    combined = pattern + '\x00' + text     # разделитель не встречается в данных
    z = z_function(combined)
    m = len(pattern)

    return [i - m - 1 for i in range(m + 1, len(combined)) if z[i] == m]


print(find_all('abracadabra', 'abra'))   # [0, 7]
print(find_all('aaaa', 'aa'))            # [0, 1, 2]

Логика простая: если z[i] == m, значит начиная с позиции i идёт полное совпадение с образцом. Разделитель нужен, чтобы совпадение не «перетекло» из образца в текст и не дало значение больше m.

Сложность — O(n + m), как у KMP.

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

Обе решают одни и те же задачи и переводятся друг в друга. Разница в удобстве.

	Z-функция	Префикс-функция (KMP)
Что считает	совпадение суффикса с началом строки	наибольший собственный префикс-суффикс
Понятность	проще	требует привыкания
Поиск подстроки	через склейку	напрямую
Автомат	сложнее построить	естественно
Онлайн-обработка	нет	да

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

Другие применения

Период строки. Минимальный период равен n - z[i] для наименьшего i, при котором i + z[i] == n. Позволяет проверить, является ли строка повторением более короткой.

def min_period(s: str) -> int:
    n = len(s)
    z = z_function(s)
    for i in range(1, n):
        if i + z[i] == n and n % i == 0:
            return i
    return n


print(min_period('abcabcabc'))   # 3
print(min_period('abcd'))        # 4

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

Сжатие строки. Найти наименьшую строку, повторением которой является исходная, — прямое следствие периода.

Поиск наибольшего общего префикса двух строк — склеиваем и смотрим z в нужной позиции.

Частые ошибки

Разделитель встречается в данных. Если взять #, а он есть в тексте, значения z могут превысить m, и алгоритм найдёт ложные вхождения. Берите заведомо отсутствующий символ — \x00 или chr(0).

min(r - i, z[i - l]) вместо min(r - i + 1, ...) или наоборот. Зависит от того, включаете вы r в блок или нет. В коде выше r — позиция за блоком, поэтому r - i без единицы. Смешение двух соглашений — источник ошибок на единицу.

Обновление l и r без проверки. Блок обновляется, только если новая правая граница больше текущей. Безусловное обновление сломает линейность.

Попытка считать z[0]. По определению это вся строка; включение его в цикл ломает логику блока.