Кореневе розбиття

Кореневе розбиття — це метод (або структура даних), що дозволяє виконувати деякі поширені операції (знаходження суми елементів підмасиву, знаходження мінімального/максимального елемента тощо) за $O(\sqrt n)$ операцій, що значно швидше за $O(n)$ у тривіальному алгоритмі.

Спершу опишемо структуру даних для одного з найпростіших застосувань цієї ідеї, потім покажемо, як узагальнити її для розв'язання деяких інших задач, і нарешті розглянемо дещо інше використання цієї ідеї: розбиття вхідних запитів на кореневі блоки.

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

• Чи влаштовує $O(\sqrt n)$ на запит заради простішої та універсальнішої реалізації? (якщо потрібне $O(\log n)$ і операція добре підтримується деревом → дерево відрізків)
• Чи виявляється операція на відрізку незручною для дерева відрізків (наприклад, потребує офлайн-обробки запитів алгоритмом Мо), де блокова структура природніша?
• Чи можна обробляти запити в офлайн-режимі, сортуючи їх за блоками для прискорення (алгоритм Мо)?

Структура даних на основі кореневого розбиття

Дано масив $a[0 \dots n-1]$. Реалізуйте структуру даних, що дозволяє знаходити суму елементів $a[l \dots r]$ для довільних $l$ і $r$ за $O(\sqrt n)$ операцій.

Опис

Основна ідея кореневого розбиття — це попередня обробка. Розіб'ємо масив $a$ на блоки завдовжки приблизно $\sqrt n$, і для кожного блока $i$ заздалегідь обчислимо суму елементів у ньому $b[i]$.

Можемо вважати, що як розмір блока, так і кількість блоків дорівнюють $\sqrt n$, заокругленому вгору:

$$s = \lceil \sqrt n \rceil$$

Тоді масив $a$ розбивається на блоки таким чином:

$$\underbrace{a[0], a[1], \dots, a[s-1]}_{\text{b[0]}}, \underbrace{a[s], \dots, a[2s-1]}_{\text{b[1]}}, \dots, \underbrace{a[(s-1) \cdot s], \dots, a[n-1]}_{\text{b[s-1]}}$$

Останній блок може мати менше елементів, ніж решта (якщо $n$ не кратне $s$), але для нашого обговорення це не важливо (бо це легко обробити). Отже, для кожного блока $k$ ми знаємо суму елементів у ньому $b[k]$:

$$b[k] = \sum\limits_{i=k\cdot s}^{\min {(n-1,(k+1)\cdot s - 1})} a[i]$$

Отже, ми обчислили значення $b[k]$ (це потребувало $O(n)$ операцій). Як вони допоможуть нам відповідати на кожен запит $[l, r]$? Зауважимо, що якщо відрізок $[l, r]$ достатньо довгий, він міститиме кілька цілих блоків, і для цих блоків ми можемо знайти суму їхніх елементів за одну операцію. Як наслідок, відрізок $[l, r]$ міститиме частини лише двох блоків, і нам доведеться обчислити суму елементів у цих частинах тривіально.

Таким чином, щоб обчислити суму елементів на відрізку $[l, r]$, нам потрібно лише підсумувати елементи двох «хвостів»: $[l\dots (k + 1)\cdot s-1]$ і $[p\cdot s\dots r]$, а також підсумувати значення $b[i]$ у всіх блоках від $k + 1$ до $p-1$:

$$\sum\limits_{i=l}^r a[i] = \sum\limits_{i=l}^{(k+1) \cdot s-1} a[i] + \sum\limits_{i=k+1}^{p-1} b[i] + \sum\limits_{i=p\cdot s}^r a[i]$$

Зауваження: коли $k = p$, тобто $l$ і $r$ належать одному блоку, формулу застосувати не можна, і суму слід обчислювати тривіально.

Цей підхід дозволяє нам значно зменшити кількість операцій. Дійсно, розмір кожного «хвоста» не перевищує довжину блока $s$, а кількість блоків у сумі не перевищує $s$. Оскільки ми обрали $s \approx \sqrt n$, загальна кількість операцій, потрібна для знаходження суми елементів на відрізку $[l, r]$, становить $O(\sqrt n)$.

Реалізація

Почнемо з найпростішої реалізації:

C++

// вхідні дані
int n;
vector<int> a (n);

// попередня обробка
int len = (int) sqrt (n + .0) + 1; // розмір блока і кількість блоків
vector<int> b (len);
for (int i=0; i<n; ++i)
    b[i / len] += a[i];

// відповіді на запити
for (;;) {
    int l, r;
  // зчитуємо вхідні дані для наступного запиту
    int sum = 0;
    for (int i=l; i<=r; )
        if (i % len == 0 && i + len - 1 <= r) {
            // якщо весь блок, що починається з i, належить [l, r]
            sum += b[i / len];
            i += len;
        }
        else {
            sum += a[i];
            ++i;
        }
}

Python

import math

# вхідні дані
n = ...
a = [...]  # масив із n елементів

# попередня обробка
length = int(math.sqrt(n)) + 1  # розмір блока і кількість блоків
b = [0] * length
for i in range(n):
    b[i // length] += a[i]  # цілі Python не переповнюються, тож сума безпечна

# відповіді на запити
while True:
    l, r = ...  # зчитуємо вхідні дані для наступного запиту
    s = 0
    i = l
    while i <= r:
        if i % length == 0 and i + length - 1 <= r:
            # якщо весь блок, що починається з i, належить [l, r]
            s += b[i // length]
            i += length
        else:
            s += a[i]
            i += 1

TypeScript

// вхідні дані
const n: number = ...;
const a: number[] = [...]; // масив із n елементів

// попередня обробка
const len = Math.floor(Math.sqrt(n)) + 1; // розмір блока і кількість блоків
const b: number[] = new Array(len).fill(0);
for (let i = 0; i < n; ++i)
    b[Math.floor(i / len)] += a[i];

// відповіді на запити
for (;;) {
    let l: number, r: number;
    // зчитуємо вхідні дані для наступного запиту
    let sum = 0;
    for (let i = l; i <= r; ) {
        if (i % len === 0 && i + len - 1 <= r) {
            // якщо весь блок, що починається з i, належить [l, r]
            sum += b[Math.floor(i / len)];
            i += len;
        } else {
            sum += a[i];
            ++i;
        }
    }
}

Go

import "math"

// вхідні дані
var n int
a := make([]int, n)

// попередня обробка
length := int(math.Sqrt(float64(n))) + 1 // розмір блока і кількість блоків
b := make([]int, length)
for i := 0; i < n; i++ {
    b[i/length] += a[i]
}

// відповіді на запити
for {
    var l, r int
    // зчитуємо вхідні дані для наступного запиту
    sum := 0
    for i := l; i <= r; {
        if i%length == 0 && i+length-1 <= r {
            // якщо весь блок, що починається з i, належить [l, r]
            sum += b[i/length]
            i += length
        } else {
            sum += a[i]
            i++
        }
    }
}

Ця реалізація має невиправдано багато операцій ділення (які значно повільніші за інші арифметичні операції). Натомість ми можемо обчислити індекси блоків $c_l$ і $c_r$, які містять індекси $l$ і $r$, і пройтися по блоках $c_l+1 \dots c_r-1$ з окремою обробкою «хвостів» у блоках $c_l$ і $c_r$. Цей підхід відповідає останній формулі з опису і робить випадок $c_l = c_r$ особливим випадком.

C++

int sum = 0;
int c_l = l / len,   c_r = r / len;
if (c_l == c_r)
    for (int i=l; i<=r; ++i)
        sum += a[i];
else {
    for (int i=l, end=(c_l+1)*len-1; i<=end; ++i)
        sum += a[i];
    for (int i=c_l+1; i<=c_r-1; ++i)
        sum += b[i];
    for (int i=c_r*len; i<=r; ++i)
        sum += a[i];
}

Python

s = 0
c_l, c_r = l // length, r // length
if c_l == c_r:
    for i in range(l, r + 1):
        s += a[i]
else:
    for i in range(l, (c_l + 1) * length):
        s += a[i]
    for i in range(c_l + 1, c_r):
        s += b[i]
    for i in range(c_r * length, r + 1):
        s += a[i]

TypeScript

let sum = 0;
const c_l = Math.floor(l / len), c_r = Math.floor(r / len);
if (c_l === c_r) {
    for (let i = l; i <= r; ++i)
        sum += a[i];
} else {
    for (let i = l, end = (c_l + 1) * len - 1; i <= end; ++i)
        sum += a[i];
    for (let i = c_l + 1; i <= c_r - 1; ++i)
        sum += b[i];
    for (let i = c_r * len; i <= r; ++i)
        sum += a[i];
}

Go

sum := 0
cL, cR := l/length, r/length
if cL == cR {
    for i := l; i <= r; i++ {
        sum += a[i]
    }
} else {
    for i, end := l, (cL+1)*length-1; i <= end; i++ {
        sum += a[i]
    }
    for i := cL + 1; i <= cR-1; i++ {
        sum += b[i]
    }
    for i := cR * length; i <= r; i++ {
        sum += a[i]
    }
}

Інші задачі

Досі ми обговорювали задачу знаходження суми елементів неперервного підмасиву. Цю задачу можна розширити, щоб дозволити оновлення окремих елементів масиву. Якщо елемент $a[i]$ змінюється, достатньо оновити значення $b[k]$ для блока, якому належить цей елемент ($k = i / s$), за одну операцію:

$$b[k] += a_{new}[i] - a_{old}[i]$$

З іншого боку, задачу знаходження суми елементів можна замінити задачею знаходження мінімального/максимального елемента підмасиву. Якщо ця задача також має враховувати оновлення окремих елементів, то оновлення значення $b[k]$ теж можливе, але воно потребуватиме проходу по всіх значеннях блока $k$ за $O(s) = O(\sqrt{n})$ операцій.

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

Інший клас задач виникає, коли нам потрібно оновлювати елементи масиву на відрізках: збільшувати наявні елементи або замінювати їх заданим значенням.

Наприклад, припустимо, що ми можемо виконувати два типи операцій над масивом: додати задане значення $\delta$ до всіх елементів масиву на відрізку $[l, r]$ або запитати значення елемента $a[i]$. Зберігатимемо в $b[k]$ значення, яке треба додати до всіх елементів блока $k$ (початково всі $b[k] = 0$). Під час кожної операції «додавання» нам потрібно додати $\delta$ до $b[k]$ для всіх блоків, що належать відрізку $[l, r]$, і додати $\delta$ до $a[i]$ для всіх елементів, що належать «хвостам» відрізка. Відповіддю на запит $i$ є просто $a[i] + b[i/s]$. Таким чином операція «додавання» має складність $O(\sqrt{n})$, а відповідь на запит — складність $O(1)$.

Нарешті, ці два класи задач можна поєднати, якщо задача вимагає виконувати і оновлення елементів на відрізку, і запити на відрізку. Обидві операції можна виконати зі складністю $O(\sqrt{n})$. Це потребуватиме двох блокових масивів $b$ і $c$: одного — для відстеження оновлень елементів, а іншого — для відстеження відповідей на запит.

Існують і інші задачі, які можна розв'язати за допомогою кореневого розбиття, наприклад, задача про підтримання множини чисел, яка дозволяла б додавати/видаляти числа, перевіряти, чи число належить множині, і знаходити $k$-те за величиною число. Щоб її розв'язати, потрібно зберігати числа у зростаючому порядку, розбитими на кілька блоків по $\sqrt{n}$ чисел у кожному. Щоразу при додаванні/видаленні числа блоки доводиться перебалансовувати, переміщуючи числа між початками і кінцями сусідніх блоків.

Алгоритм Мо

Подібну ідею, що ґрунтується на кореневому розбитті, можна використати для відповіді на запити на відрізках ($Q$) в офлайн-режимі за $O((N+Q)\sqrt{N})$. Це може звучати значно гірше за методи з попереднього розділу, оскільки це дещо гірша складність, ніж була раніше, і не дозволяє оновлювати значення між двома запитами. Але в багатьох ситуаціях цей метод має переваги. Під час звичайного кореневого розбиття нам доводиться заздалегідь обчислювати відповіді для кожного блока і об'єднувати їх під час відповіді на запити. У деяких задачах цей крок об'єднання може бути доволі проблематичним. Наприклад, коли кожен запит просить знайти моду свого відрізка (число, що зустрічається найчастіше). Для цього кожному блоку довелося б зберігати кількість входжень кожного числа в ньому в якійсь структурі даних, і ми вже не можемо виконувати крок об'єднання достатньо швидко. Алгоритм Мо використовує цілком інший підхід, що дозволяє відповідати на запити такого роду швидко, бо він підтримує лише одну структуру даних, а єдині операції з нею прості й швидкі.

Ідея полягає в тому, щоб відповідати на запити в особливому порядку, що ґрунтується на індексах. Спершу ми відповімо на всі запити, у яких лівий індекс лежить у блоці 0, потім — на всі запити, у яких лівий індекс лежить у блоці 1, і так далі. А ще нам доведеться відповідати на запити одного блока в особливому порядку, а саме відсортованими за правим індексом запитів.

Як уже сказано, ми використовуватимемо єдину структуру даних. Ця структура даних зберігатиме інформацію про відрізок. Спочатку цей відрізок буде порожнім. Коли ми хочемо відповісти на наступний запит (в особливому порядку), ми просто розширюємо або звужуємо відрізок, додаючи/видаляючи елементи з обох боків поточного відрізка, доки не перетворимо його на відрізок запиту. У такий спосіб нам потрібно лише додати або видалити один елемент за раз, що має бути доволі простими операціями в нашій структурі даних.

Оскільки ми змінюємо порядок відповіді на запити, це можливо лише тоді, коли нам дозволено відповідати на запити в офлайн-режимі.

Реалізація

В алгоритмі Мо ми використовуємо дві функції — для додавання індекса і для видалення індекса з відрізка, який ми наразі підтримуємо.

C++

void remove(idx);  // TODO: видалити значення за індексом idx зі структури даних
void add(idx);     // TODO: додати значення за індексом idx до структури даних
int get_answer();  // TODO: видобути поточну відповідь зі структури даних

int block_size;

struct Query {
    int l, r, idx;
    bool operator<(Query other) const
    {
        return make_pair(l / block_size, r) <
               make_pair(other.l / block_size, other.r);
    }
};

vector<int> mo_s_algorithm(vector<Query> queries) {
    vector<int> answers(queries.size());
    sort(queries.begin(), queries.end());

    // TODO: ініціалізувати структуру даних

    int cur_l = 0;
    int cur_r = -1;
    // інваріант: структура даних завжди відображає відрізок [cur_l, cur_r]
    for (Query q : queries) {
        while (cur_l > q.l) {
            cur_l--;
            add(cur_l);
        }
        while (cur_r < q.r) {
            cur_r++;
            add(cur_r);
        }
        while (cur_l < q.l) {
            remove(cur_l);
            cur_l++;
        }
        while (cur_r > q.r) {
            remove(cur_r);
            cur_r--;
        }
        answers[q.idx] = get_answer();
    }
    return answers;
}

Python

def remove(idx):  # TODO: видалити значення за індексом idx зі структури даних
    ...

def add(idx):     # TODO: додати значення за індексом idx до структури даних
    ...

def get_answer():  # TODO: видобути поточну відповідь зі структури даних
    ...

block_size = ...

# кожен запит — це кортеж (l, r, idx)
def mos_algorithm(queries):
    answers = [0] * len(queries)
    # сортуємо за (блок лівого індекса, правий індекс)
    queries.sort(key=lambda q: (q[0] // block_size, q[1]))

    # TODO: ініціалізувати структуру даних

    cur_l = 0
    cur_r = -1
    # інваріант: структура даних завжди відображає відрізок [cur_l, cur_r]
    for l, r, idx in queries:
        while cur_l > l:
            cur_l -= 1
            add(cur_l)
        while cur_r < r:
            cur_r += 1
            add(cur_r)
        while cur_l < l:
            remove(cur_l)
            cur_l += 1
        while cur_r > r:
            remove(cur_r)
            cur_r -= 1
        answers[idx] = get_answer()
    return answers

TypeScript

function remove(idx: number): void {} // TODO: видалити значення за індексом idx зі структури даних
function add(idx: number): void {}    // TODO: додати значення за індексом idx до структури даних
function getAnswer(): number { return 0; } // TODO: видобути поточну відповідь зі структури даних

const blockSize: number = ...;

interface Query {
    l: number;
    r: number;
    idx: number;
}

function mosAlgorithm(queries: Query[]): number[] {
    const answers: number[] = new Array(queries.length).fill(0);
    // сортуємо за (блок лівого індекса, правий індекс)
    queries.sort((a, c) => {
        const ba = Math.floor(a.l / blockSize), bc = Math.floor(c.l / blockSize);
        return ba !== bc ? ba - bc : a.r - c.r;
    });

    // TODO: ініціалізувати структуру даних

    let cur_l = 0;
    let cur_r = -1;
    // інваріант: структура даних завжди відображає відрізок [cur_l, cur_r]
    for (const q of queries) {
        while (cur_l > q.l) {
            cur_l--;
            add(cur_l);
        }
        while (cur_r < q.r) {
            cur_r++;
            add(cur_r);
        }
        while (cur_l < q.l) {
            remove(cur_l);
            cur_l++;
        }
        while (cur_r > q.r) {
            remove(cur_r);
            cur_r--;
        }
        answers[q.idx] = getAnswer();
    }
    return answers;
}

Go

func remove(idx int) {}      // TODO: видалити значення за індексом idx зі структури даних
func add(idx int) {}         // TODO: додати значення за індексом idx до структури даних
func getAnswer() int { return 0 } // TODO: видобути поточну відповідь зі структури даних

var blockSize int

type Query struct {
    l, r, idx int
}

func mosAlgorithm(queries []Query) []int {
    answers := make([]int, len(queries))
    // сортуємо за (блок лівого індекса, правий індекс)
    sort.Slice(queries, func(i, j int) bool {
        bi, bj := queries[i].l/blockSize, queries[j].l/blockSize
        if bi != bj {
            return bi < bj
        }
        return queries[i].r < queries[j].r
    })

    // TODO: ініціалізувати структуру даних

    curL := 0
    curR := -1
    // інваріант: структура даних завжди відображає відрізок [curL, curR]
    for _, q := range queries {
        for curL > q.l {
            curL--
            add(curL)
        }
        for curR < q.r {
            curR++
            add(curR)
        }
        for curL < q.l {
            remove(curL)
            curL++
        }
        for curR > q.r {
            remove(curR)
            curR--
        }
        answers[q.idx] = getAnswer()
    }
    return answers
}

Залежно від задачі ми можемо використовувати різні структури даних і відповідно змінювати функції add/remove/get_answer. Наприклад, якщо нас просять відповідати на запити суми на відрізку, то ми використовуємо просте ціле число як структуру даних, яке на початку дорівнює $0$. Функція add просто додаватиме значення позиції і потім оновлюватиме змінну-відповідь. З іншого боку, функція remove віднімає значення в позиції і потім оновлює змінну-відповідь. А get_answer просто повертає це ціле число.

Для відповіді на запити моди ми можемо використати бінарне дерево пошуку (наприклад, map<int, int>) для зберігання того, як часто кожне число зустрічається в поточному відрізку, і друге бінарне дерево пошуку (наприклад, set<pair<int, int>>) для зберігання кількостей чисел (наприклад, у вигляді пар кількість-число) у порядку. Метод add видаляє поточне число з другого BST, збільшує кількість у першому і вставляє число назад у друге. remove робить те саме, лише зменшує кількість. А get_answer просто дивиться на друге дерево і повертає найкраще значення за $O(1)$.

Складність

Сортування всіх запитів займе $O(Q \log Q)$.

А як щодо інших операцій? Скільки разів будуть викликані add і remove?

Скажімо, розмір блока дорівнює $S$.

Якщо ми дивимося лише на всі запити, у яких лівий індекс лежить в одному й тому самому блоці, то ці запити відсортовані за правим індексом. Тому ми викличемо add(cur_r) і remove(cur_r) лише $O(N)$ разів для всіх цих запитів разом. Це дає $O(\frac{N}{S} N)$ викликів для всіх блоків.

Значення cur_l може змінитися щонайбільше на $O(S)$ між двома запитами. Тому маємо додаткові $O(S Q)$ викликів add(cur_l) і remove(cur_l).

Для $S \approx \sqrt{N}$ це дає загалом $O((N + Q) \sqrt{N})$ операцій. Отже, складність становить $O((N+Q)F\sqrt{N})$, де $O(F)$ — це складність функцій add і remove.

Поради щодо покращення швидкодії

• Розмір блока, що дорівнює точно $\sqrt{N}$, не завжди дає найкращу швидкодію. Наприклад, якщо $\sqrt{N}=750$, то може статися, що розмір блока $700$ або $800$ працюватиме краще. Що важливіше, не обчислюйте розмір блока під час виконання — зробіть його const. Ділення на константи добре оптимізується компіляторами.
• У непарних блоках сортуйте правий індекс за зростанням, а в парних блоках — за спаданням. Це мінімізує рух правого вказівника, бо звичайне сортування переміщуватиме правий вказівник з кінця назад на початок на старті кожного блока. З покращеною версією таке скидання більше не потрібне.

C++

bool cmp(pair<int, int> p, pair<int, int> q) {
    if (p.first / BLOCK_SIZE != q.first / BLOCK_SIZE)
        return p < q;
    return (p.first / BLOCK_SIZE & 1) ? (p.second < q.second) : (p.second > q.second);
}

Python

import functools

def cmp(p, q):
    # p, q — це пари (first, second)
    bp, bq = p[0] // BLOCK_SIZE, q[0] // BLOCK_SIZE
    if bp != bq:
        return -1 if p < q else 1
    if bp & 1:
        return -1 if p[1] < q[1] else (1 if p[1] > q[1] else 0)
    return -1 if p[1] > q[1] else (1 if p[1] < q[1] else 0)

# використання: queries.sort(key=functools.cmp_to_key(cmp))

TypeScript

// p, q — це пари [first, second]; повертає < 0, якщо p йде раніше за q
function cmp(p: [number, number], q: [number, number]): number {
    const bp = Math.floor(p[0] / BLOCK_SIZE), bq = Math.floor(q[0] / BLOCK_SIZE);
    if (bp !== bq)
        return p[0] !== q[0] ? p[0] - q[0] : p[1] - q[1];
    return (bp & 1) ? (p[1] - q[1]) : (q[1] - p[1]);
}

Go

// p, q — це пари [2]int{first, second}; повертає true, якщо p йде раніше за q
func cmp(p, q [2]int) bool {
    bp, bq := p[0]/BLOCK_SIZE, q[0]/BLOCK_SIZE
    if bp != bq {
        if p[0] != q[0] {
            return p[0] < q[0]
        }
        return p[1] < q[1]
    }
    if bp&1 == 1 {
        return p[1] < q[1]
    }
    return p[1] > q[1]
}

Про ще швидший підхід до сортування ви можете прочитати тут.

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

• Codeforces - Kuriyama Mirai's Stones
• UVA - 12003 - Array Transformer
• UVA - 11990 Dynamic Inversion
• SPOJ - Give Away
• Codeforces - Till I Collapse
• Codeforces - Destiny
• Codeforces - Holes
• Codeforces - XOR and Favorite Number
• Codeforces - Powerful array
• SPOJ - DQUERY
• Codeforces - Robin Hood Archery

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