Treap | kz-maxx

Түсіндіру

Treap — бұл декартово дерево. Ол екі идеяны бірге қолданады:

Binary Search Tree — key бойынша
Heap — кездейсоқ prior бойынша

key бойынша: сол жақта кіші мәндер, оң жақта үлкен мәндер тұрады. prior бойынша: parent prior балаларының prior мәнінен үлкен болады.

key: left < root < right
prior: root.prior > child.prior

Приоритет кездейсоқ болғандықтан, дерево көбіне теңгерілген болады. Сондықтан негізгі операциялар орташа O(log n) уақытта жұмыс істейді.

Массивте не сақтаймыз?

struct Node {
    int key, prior;
    int l, r;
    int cnt;
} tr[N];

int root;

key — вершинадағы мән
prior — кездейсоқ приоритет
l — сол жақ баланың индексі
r — оң жақ баланың индексі
cnt — осы поддеревоның размері
root — барлық деревоның корені

0 индексі пустая вершина дегенді білдіреді. Яғни бұл nullptr сияқты, бірақ массивте.

Негізгі операциялар

split

split(v, x) дерево v-ны екі бөлікке бөледі:

бірінші бөлік — барлық мәндер <= x
екінші бөлік — барлық мәндер > x

split(root, 10) → {дерево <= 10, дерево > 10}

merge

merge(a, b) екі деревоны бір деревоға қосады.

Маңызды: a ішіндегі барлық мәндер b ішіндегі барлық мәндерден кіші болуы керек.

merge(a, b) → бір дерево

Реализацияға арналған есеп мысалы

Бізге сандар жиынын сақтау керек және бірнеше команда орындау керек.

Команда	Не істейді
insert x	x санын қосады
erase x	x санын өшіреді
find x	x бар ма, соны тексереді
kth k	өсу ретімен k-шы элементті табады

Input

10
insert 10
insert 5
insert 20
find 10
find 7
kth 2
erase 10
find 10
kth 2
find 20

Output

YES
NO
10
NO
20
YES

Толық реализация

#include <bits/stdc++.h>
using namespace std;

const int N = 100010;

struct Node {
    int key, prior;
    int l, r;
    int cnt;
} tr[N];

int root;
int idx;

int sizeOf(int v) {
    if (v == 0) {
        return 0;
    }

    return tr[v].cnt;
}

void pull(int v) {
    if (v == 0) {
        return;
    }

    tr[v].cnt = 1 + sizeOf(tr[v].l) + sizeOf(tr[v].r);
}

int newNode(int x) {
    idx++;

    tr[idx].key = x;
    tr[idx].prior = rand();
    tr[idx].l = 0;
    tr[idx].r = 0;
    tr[idx].cnt = 1;

    return idx;
}

pair<int, int> split(int v, int x) {
    if (v == 0) {
        return {0, 0};
    }

    if (tr[v].key <= x) {
        pair<int, int> p = split(tr[v].r, x);

        tr[v].r = p.first;
        pull(v);

        return {v, p.second};
    } else {
        pair<int, int> p = split(tr[v].l, x);

        tr[v].l = p.second;
        pull(v);

        return {p.first, v};
    }
}

int merge(int a, int b) {
    if (a == 0) {
        return b;
    }

    if (b == 0) {
        return a;
    }

    if (tr[a].prior > tr[b].prior) {
        tr[a].r = merge(tr[a].r, b);
        pull(a);

        return a;
    } else {
        tr[b].l = merge(a, tr[b].l);
        pull(b);

        return b;
    }
}

bool findKey(int v, int x) {
    if (v == 0) {
        return false;
    }

    if (tr[v].key == x) {
        return true;
    }

    if (x < tr[v].key) {
        return findKey(tr[v].l, x);
    } else {
        return findKey(tr[v].r, x);
    }
}

void insertKey(int x) {
    if (findKey(root, x)) {
        return;
    }

    pair<int, int> p = split(root, x);

    int leftTree = p.first;
    int rightTree = p.second;
    int node = newNode(x);

    root = merge(merge(leftTree, node), rightTree);
}

void eraseKey(int &v, int x) {
    if (v == 0) {
        return;
    }

    if (tr[v].key == x) {
        v = merge(tr[v].l, tr[v].r);
        return;
    }

    if (x < tr[v].key) {
        eraseKey(tr[v].l, x);
    } else {
        eraseKey(tr[v].r, x);
    }

    pull(v);
}

int kth(int v, int k) {
    int leftSize = sizeOf(tr[v].l);

    if (k == leftSize + 1) {
        return tr[v].key;
    }

    if (k <= leftSize) {
        return kth(tr[v].l, k);
    }

    return kth(tr[v].r, k - leftSize - 1);
}

void printTree(int v) {
    if (v == 0) {
        return;
    }

    printTree(tr[v].l);
    cout << tr[v].key << ' ';
    printTree(tr[v].r);
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    srand(time(0));

    int q;
    cin >> q;

    while (q--) {
        string op;
        int x;

        cin >> op >> x;

        if (op == "insert") {
            insertKey(x);
        } else if (op == "erase") {
            eraseKey(root, x);
        } else if (op == "find") {
            if (findKey(root, x)) {
                cout << "YES\n";
            } else {
                cout << "NO\n";
            }
        } else if (op == "kth") {
            if (x >= 1 && x <= sizeOf(root)) {
                cout << kth(root, x) << '\n';
            }
        }
    }

    return 0;
}

Функциялар түсіндірмесі

1. sizeOf

int sizeOf(int v) {
    if (v == 0) {
        return 0;
    }

    return tr[v].cnt;
}

Бұл функция поддерево размерін қайтарады. Егер вершина жоқ болса, размер 0 болады.

sizeOf(0) = 0
sizeOf(v) = tr[v].cnt

2. pull

void pull(int v) {
    tr[v].cnt = 1 + sizeOf(tr[v].l) + sizeOf(tr[v].r);
}

Бұл функция v вершинасының cnt мәнін қайта есептейді.

сол жақта 2 вершина
оң жақта 3 вершина
cnt = 1 + 2 + 3 = 6

3. newNode

int newNode(int x)

Жаңа вершина жасайды және оның индексін қайтарады.

newNode(10)
key = 10
l = 0, r = 0
cnt = 1

4. split

pair<int, int> split(int v, int x)

Деревоны екі бөлікке бөледі:

бірінші бөлік: мәндер <= x
екінші бөлік: мәндер > x

дерево: 3 5 8 10 15
split(root, 8)
leftTree: 3 5 8
rightTree: 10 15

Егер қазіргі key <= x болса, онда бұл вершина сол жақ бөлікке кіреді. Сондықтан оның оң жақ баласын бөлеміз.

if (tr[v].key <= x) {
    pair<int, int> p = split(tr[v].r, x);

    tr[v].r = p.first;
    pull(v);

    return {v, p.second};
}

Егер қазіргі key > x болса, онда бұл вершина оң жақ бөлікке кіреді. Сондықтан оның сол жақ баласын бөлеміз.

else {
    pair<int, int> p = split(tr[v].l, x);

    tr[v].l = p.second;
    pull(v);

    return {p.first, v};
}

5. merge

int merge(int a, int b)

Екі деревоны бір деревоға қосады. a ішіндегі барлық мәндер b ішіндегі мәндерден кіші болуы керек.

a: 1 3 5
b: 8 10 12
merge(a, b): 1 3 5 8 10 12

Қай вершинада prior үлкен болса, сол корень болады.

if (tr[a].prior > tr[b].prior) {
    tr[a].r = merge(tr[a].r, b);
    pull(a);

    return a;
}

Егер b приоритеті үлкен болса, онда b корень болады.

else {
    tr[b].l = merge(a, tr[b].l);
    pull(b);

    return b;
}

6. findKey

bool findKey(int v, int x)

x санын іздейді. Бұл кәдімгі binary search tree сияқты жұмыс істейді.

егер x кіші болса → солға барамыз
егер x үлкен болса → оңға барамыз
егер x == key болса → таптық

7. insertKey

void insertKey(int x)

x санын деревоға қосады.

бұрын: 3 5 10
insert 7
split(root, 7): 3 5 және 10
кейін: merge(3 5, 7, 10)
болды: 3 5 7 10

pair<int, int> p = split(root, x);

int leftTree = p.first;
int rightTree = p.second;
int node = newNode(x);

root = merge(merge(leftTree, node), rightTree);

8. eraseKey

void eraseKey(int &v, int x)

x санын дерево ішінен өшіреді. Керек вершина табылса, оны сол және оң балаларының merge нәтижесімен ауыстырамыз.

бұрын: left ← x → right
кейін: merge(left, right)

if (tr[v].key == x) {
    v = merge(tr[v].l, tr[v].r);
    return;
}

9. kth

int kth(int v, int k)

Өсу ретімен k-шы элементті табады. Ол үшін сол жақ поддерево размері қолданылады.

дерево: 3 5 7 10 20
kth(1) = 3
kth(3) = 7
kth(5) = 20

Егер сол жақта leftSize элемент болса:

егер k == leftSize + 1, жауап — қазіргі вершина
егер k <= leftSize, сол жаққа барамыз
әйтпесе оң жаққа барамыз және k мәнін азайтамыз

10. printTree

void printTree(int v)

Деревоны өсу ретімен шығарады. Себебі алдымен сол жақ, кейін root, соңында оң жақ шығарылады.

дерево ішінде: 10, 3, 20, 5
printTree шығарады: 3 5 10 20

Асимптотика

Операция	Орташа уақыт
insert	O(log n)
erase	O(log n)
find	O(log n)
kth	O(log n)
split / merge	O(log n)