Z-функція та її обчислення

Нехай задано рядок $s$ довжини $n$. Z-функція цього рядка — це масив довжини $n$, де $i$-й елемент дорівнює найбільшій кількості символів, починаючи з позиції $i$, які збігаються з першими символами $s$.

Іншими словами, $z[i]$ — це довжина найдовшого рядка, який одночасно є префіксом $s$ і префіксом суфікса $s$, що починається з позиції $i$.

Зауваження. У цій статті, щоб уникнути двозначності, ми вважаємо, що індексація починається з $0$; тобто перший символ $s$ має індекс $0$, а останній — індекс $n-1$.

Перший елемент Z-функції, $z[0]$, загалом не є чітко визначеним. У цій статті ми вважатимемо його рівним нулю (хоча це нічого не змінює в реалізації алгоритму).

У цій статті наведено алгоритм обчислення Z-функції за час $O(n)$, а також різні її застосування.

Коли підходить цей алгоритм?

• Шукаєте один взірець у тексті чи аналізуєте збіги з префіксом одного рядка? (якщо взірців багато одночасно → Ахо-Корасік)
• Достатньо точного збігу, а не порівняння довільних підрядків за $O(1)$? (якщо потрібні саме такі порівняння → Хешування рядків)
• Вам зручніше міркувати у термінах «довжина збігу з префіксом», ніж бордерів? (якщо зручніша префікс-функція → КМП)

Приклади

Наприклад, ось значення Z-функції, обчислені для різних рядків:

• "aaaaa" — $[0, 4, 3, 2, 1]$
• "aaabaab" — $[0, 2, 1, 0, 2, 1, 0]$
• "abacaba" — $[0, 0, 1, 0, 3, 0, 1]$

Тривіальний алгоритм

Формальне означення можна записати у вигляді такої елементарної реалізації за час $O(n^2)$.

C++

vector<int> z_function_trivial(string s) {
	int n = s.size();
	vector<int> z(n);
	for (int i = 1; i < n; i++) {
		while (i + z[i] < n && s[z[i]] == s[i + z[i]]) {
			z[i]++;
		}
	}
	return z;
}

Python

def z_function_trivial(s: str) -> list[int]:
    n = len(s)
    z = [0] * n
    for i in range(1, n):
        # збільшуємо z[i], доки символи збігаються з префіксом
        while i + z[i] < n and s[z[i]] == s[i + z[i]]:
            z[i] += 1
    return z

TypeScript

function zFunctionTrivial(s: string): number[] {
    const n = s.length;
    const z = new Array<number>(n).fill(0);
    for (let i = 1; i < n; i++) {
        // збільшуємо z[i], доки символи збігаються з префіксом
        while (i + z[i] < n && s[z[i]] === s[i + z[i]]) {
            z[i]++;
        }
    }
    return z;
}

Go

func zFunctionTrivial(s string) []int {
    n := len(s)
    z := make([]int, n)
    for i := 1; i < n; i++ {
        // збільшуємо z[i], доки символи збігаються з префіксом
        for i+z[i] < n && s[z[i]] == s[i+z[i]] {
            z[i]++
        }
    }
    return z
}

Ми просто проходимо по кожній позиції $i$ та оновлюємо $z[i]$ для кожної з них, починаючи з $z[i] = 0$ і збільшуючи його, доки не знайдемо розбіжність (і доки не досягнемо кінця рядка).

Звісно, це неефективна реалізація. Тепер ми покажемо побудову ефективної реалізації.

Ефективний алгоритм обчислення Z-функції

Щоб отримати ефективний алгоритм, ми обчислюватимемо значення $z[i]$ по черзі від $i = 1$ до $n - 1$, але водночас, обчислюючи нове значення, намагатимемося якнайкраще скористатися вже обчисленими значеннями.

Для стислості назвемо Z-блоками (англ. segment match) ті підрядки, які збігаються з префіксом $s$. Наприклад, значення шуканої Z-функції $z[i]$ — це довжина Z-блоку, що починається з позиції $i$ (і закінчується на позиції $i + z[i] - 1$).

Для цього ми зберігатимемо індекси $[l, r)$ найправішого Z-блоку. Тобто серед усіх знайдених блоків ми зберігаємо той, що закінчується найправіше. Певною мірою індекс $r$ можна розглядати як «межу», до якої наш рядок $s$ був переглянутий алгоритмом; усе, що далі цієї точки, ще не відоме.

Тоді, якщо поточний індекс (для якого нам треба обчислити наступне значення Z-функції) дорівнює $i$, маємо один із двох варіантів:

• $i \geq r$ — поточна позиція поза тим, що ми вже обробили.

  Тоді ми обчислимо $z[i]$ за допомогою тривіального алгоритму (тобто просто порівнюючи значення одне за одним). Зауважимо, що зрештою, якщо $z[i] > 0$, нам доведеться оновити індекси найправішого блоку, бо гарантовано новий $r = i + z[i]$ кращий за попередній $r$.

• $i < r$ — поточна позиція всередині поточного Z-блоку $[l, r)$.

  Тоді ми можемо скористатися вже обчисленими значеннями Z-функції, щоб «ініціалізувати» значення $z[i]$ чимось (це напевно краще, ніж «починати з нуля»), можливо, навіть якимось великим числом.

  Для цього зауважимо, що підрядки $s[l \dots r)$ і $s[0 \dots r-l)$ збігаються. Це означає, що як початкове наближення для $z[i]$ ми можемо взяти значення, вже обчислене для відповідного блоку $s[0 \dots r-l)$, а саме $z[i-l]$.

  Однак значення $z[i-l]$ може виявитися завеликим: застосоване до позиції $i$, воно може перевищити індекс $r$. Це неприпустимо, бо ми нічого не знаємо про символи праворуч від $r$: вони можуть відрізнятися від потрібних.

  Ось приклад подібної ситуації:

  $$s = "aaaabaa"$$

  Коли ми доходимо до останньої позиції ($i = 6$), поточним Z-блоком буде $[5, 7)$. Позиція $6$ тоді відповідатиме позиції $6 - 5 = 1$, для якої значення Z-функції дорівнює $z[1] = 3$. Очевидно, ми не можемо ініціалізувати $z[6]$ значенням $3$, це було б цілком неправильно. Максимальне значення, яким ми могли б його ініціалізувати, — це $1$, бо це найбільше значення, яке не виводить нас за межі індексу $r$ Z-блоку $[l, r)$.

  Отже, як початкове наближення для $z[i]$ ми можемо безпечно взяти:

  $$z_0[i] = \min(r - i,\; z[i-l])$$

  Після ініціалізації $z[i]$ значенням $z_0[i]$ ми намагаємося збільшити $z[i]$, запускаючи тривіальний алгоритм, бо загалом, після межі $r$, ми не можемо знати, чи продовжуватиметься збіг, чи ні.

Отже, увесь алгоритм розбивається на два випадки, які відрізняються лише початковим значенням $z[i]$: у першому випадку воно вважається рівним нулю, у другому визначається за вже обчисленими значеннями (за наведеною вище формулою). Після цього обидві гілки цього алгоритму зводяться до реалізації тривіального алгоритму, який стартує одразу після задання початкового значення.

Алгоритм виявляється дуже простим. Попри те, що на кожній ітерації запускається тривіальний алгоритм, ми досягли значного прогресу, отримавши алгоритм, який працює за лінійний час. Далі ми доведемо, що час роботи лінійний.

Реалізація

Реалізація виявляється доволі лаконічною:

C++

vector<int> z_function(string s) {
    int n = s.size();
    vector<int> z(n);
    int l = 0, r = 0;
    for(int i = 1; i < n; i++) {
        if(i < r) {
            z[i] = min(r - i, z[i - l]);
        }
        while(i + z[i] < n && s[z[i]] == s[i + z[i]]) {
            z[i]++;
        }
        if(i + z[i] > r) {
            l = i;
            r = i + z[i];
        }
    }
    return z;
}

Python

def z_function(s: str) -> list[int]:
    n = len(s)
    z = [0] * n
    l, r = 0, 0  # межі [l, r) найправішого Z-блоку
    for i in range(1, n):
        if i < r:
            # початкове наближення з вже обчисленого значення
            z[i] = min(r - i, z[i - l])
        # намагаємося збільшити z[i] тривіальним алгоритмом
        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

TypeScript

function zFunction(s: string): number[] {
    const n = s.length;
    const z = new Array<number>(n).fill(0);
    let l = 0, r = 0; // межі [l, r) найправішого Z-блоку
    for (let i = 1; i < n; i++) {
        if (i < r) {
            // початкове наближення з вже обчисленого значення
            z[i] = Math.min(r - i, z[i - l]);
        }
        // намагаємося збільшити z[i] тривіальним алгоритмом
        while (i + z[i] < n && s[z[i]] === s[i + z[i]]) {
            z[i]++;
        }
        if (i + z[i] > r) {
            l = i;
            r = i + z[i];
        }
    }
    return z;
}

Go

func zFunction(s string) []int {
    n := len(s)
    z := make([]int, n)
    l, r := 0, 0 // межі [l, r) найправішого Z-блоку
    for i := 1; i < n; i++ {
        if i < r {
            // початкове наближення з вже обчисленого значення
            z[i] = min(r-i, z[i-l])
        }
        // намагаємося збільшити z[i] тривіальним алгоритмом
        for i+z[i] < n && s[z[i]] == s[i+z[i]] {
            z[i]++
        }
        if i+z[i] > r {
            l, r = i, i+z[i]
        }
    }
    return z
}

func min(a, b int) int {
    if a < b {
        return a
    }
    return b
}

Коментарі до цієї реалізації

Весь розв'язок подано у вигляді функції, яка повертає масив довжини $n$ — Z-функцію рядка $s$.

Масив $z$ початково заповнений нулями. Поточним найправішим Z-блоком вважається $[0; 0)$ (тобто навмисно малий блок, який не містить жодного $i$).

Усередині циклу для $i = 1 \dots n - 1$ ми спершу визначаємо початкове значення $z[i]$ — воно або залишиться нулем, або буде обчислене за наведеною вище формулою.

Після цього тривіальний алгоритм намагається збільшити значення $z[i]$ якомога більше.

Зрештою, якщо це потрібно (тобто якщо $i + z[i] > r$), ми оновлюємо найправіший Z-блок $[l, r)$.

Асимптотична поведінка алгоритму

Ми доведемо, що наведений вище алгоритм має час роботи, лінійний за довжиною рядка, — тобто $O(n)$.

Доведення дуже просте.

Нас цікавить вкладений цикл while, оскільки все інше — це лише набір сталих операцій, які в сумі дають $O(n)$.

Ми покажемо, що кожна ітерація циклу while збільшує праву межу $r$ Z-блоку.

Для цього розглянемо обидві гілки алгоритму:

• $i \geq r$

  У цьому випадку або цикл while не зробить жодної ітерації (якщо $s[0] \ne s[i]$), або зробить кілька ітерацій, починаючи з позиції $i$, щоразу зсуваючись на один символ праворуч. Після цього права межа $r$ обов'язково оновиться.

  Отже, ми з'ясували, що коли $i \geq r$, кожна ітерація циклу while збільшує значення нового індексу $r$.

• $i < r$

  У цьому випадку ми ініціалізуємо $z[i]$ певним значенням $z_0$, заданим наведеною вище формулою. Порівняймо це початкове значення $z_0$ зі значенням $r - i$. Матимемо три випадки:

  • $z_0 < r - i$

    Доведемо, що в цьому випадку не відбудеться жодної ітерації циклу while.

    Це легко довести, наприклад, від супротивного: якби цикл while зробив хоча б одну ітерацію, це означало б, що початкове наближення $z[i] = z_0$ було неточним (меншим за фактичну довжину збігу). Але оскільки $s[l \dots r)$ і $s[0 \dots r-l)$ однакові, це означало б, що $z[i-l]$ містить хибне значення (менше, ніж мало б бути).

    Отже, оскільки $z[i-l]$ правильне і менше за $r - i$, звідси випливає, що це значення збігається з шуканим значенням $z[i]$.

  • $z_0 = r - i$

    У цьому випадку цикл while може зробити кілька ітерацій, але кожна з них приведе до збільшення значення індексу $r$, бо ми почнемо порівнювати з $s[r]$, який вийде за межі інтервалу $[l, r)$.

  • $z_0 > r - i$

    Цей варіант неможливий за означенням $z_0$.

Отже, ми довели, що кожна ітерація внутрішнього циклу просуває вказівник $r$ праворуч. Оскільки $r$ не може бути більшим за $n-1$, це означає, що внутрішній цикл зробить не більш ніж $n-1$ ітерацій.

Оскільки решта алгоритму очевидно працює за $O(n)$, ми довели, що весь алгоритм обчислення Z-функції працює за лінійний час.

Застосування

Тепер розглянемо деякі способи використання Z-функції для конкретних задач.

Ці застосування значною мірою подібні до застосувань префікс-функції.

Пошук підрядка

Щоб уникнути плутанини, назвемо $t$ рядком тексту, а $p$ — взірцем. Задача така: знайти всі входження взірця $p$ у текст $t$.

Щоб розв'язати цю задачу, побудуємо новий рядок $s = p + \diamond + t$, тобто конкатенуємо $p$ і $t$, але посередині поставимо символ-роздільник $\diamond$ (ми оберемо $\diamond$ так, щоб він напевно не зустрічався ніде в рядках $p$ чи $t$).

Обчислимо Z-функцію для $s$. Тоді для будь-якого $i$ з інтервалу $[0; \; \operatorname{length}(t) - 1]$ ми розглянемо відповідне значення $k = z[i + \operatorname{length}(p) + 1]$. Якщо $k$ дорівнює $\operatorname{length}(p)$, то ми знаємо, що є одне входження $p$ у $i$-й позиції $t$; інакше входження $p$ у $i$-й позиції $t$ немає.

Час роботи (і споживання пам'яті) становить $O(\operatorname{length}(t) + \operatorname{length}(p))$.

Кількість різних підрядків у рядку

Заданий рядок $s$ довжини $n$. Підрахувати кількість різних підрядків $s$.

Ми розв'яжемо цю задачу ітеративно. Тобто: знаючи поточну кількість різних підрядків, перерахуємо цю кількість після додавання в кінець $s$ одного символу.

Отже, нехай $k$ — поточна кількість різних підрядків $s$. Додамо новий символ $c$ до $s$. Очевидно, що можуть з'явитися деякі нові підрядки, що закінчуються цим новим символом $c$ (а саме всі ті рядки, які закінчуються цим символом і яких ми ще не зустрічали).

Візьмемо рядок $t = s + c$ і обернемо його (запишемо його символи у зворотному порядку). Тепер наше завдання — підрахувати, скільки префіксів $t$ не зустрічаються ніде більше в $t$. Обчислимо Z-функцію $t$ і знайдемо її максимальне значення $z_{max}$. Очевидно, що префікс $t$ довжини $z_{max}$ зустрічається також десь усередині $t$. Зрозуміло, що коротші префікси також зустрічаються.

Отже, ми з'ясували, що кількість нових підрядків, які з'являються при додаванні символу $c$ до $s$, дорівнює $\operatorname{length}(t) - z_{max}$.

Відповідно, час роботи цього розв'язку становить $O(n^2)$ для рядка довжини $n$.

Варто зазначити, що точно так само ми можемо перераховувати, все ще за час $O(n)$, кількість різних підрядків при додаванні символу на початок рядка, а також при його вилученні (з кінця або з початку).

Стиснення рядка

Заданий рядок $s$ довжини $n$. Знайти його найкоротше «стиснене» представлення, тобто: знайти рядок $t$ найкоротшої довжини такий, що $s$ можна подати як конкатенацію однієї або кількох копій $t$.

Розв'язок такий: обчислити Z-функцію $s$, пройти по всіх $i$ таких, що $i$ ділить $n$. Зупинитися на першому $i$ такому, що $i + z[i] = n$. Тоді рядок $s$ можна стиснути до довжини $i$.

Доведення цього факту таке саме, як і для розв'язку, що використовує префікс-функцію.

Задачі для практики

• CSES - Finding Borders
• eolymp - Blocks of string
• Codeforces - Password [Difficulty: Easy]
• UVA # 455 "Periodic Strings" [Difficulty: Medium]
• UVA # 11022 "String Factoring" [Difficulty: Medium]
• UVa 11475 - Extend to Palindrome
• LA 6439 - Pasti Pas!
• Codechef - Chef and Strings
• Codeforces - Prefixes and Suffixes

Відеоматеріали