← Назад к структурам данных

АВЛ-дерево

Сбалансированное двоичное дерево поиска с поддержкой порядка

🐍 Python
def fh(i):
    if i == -1: return 0
    return tree[i][3]

def fk(i):
    if i == -1: return 0
    return tree[i][5]

def upd(i):
    if i == -1: return
    l = tree[i][1]
    r = tree[i][2]
    tree[i][3] = 1 + max(fh(l), fh(r))
    tree[i][5] = tree[i][4] + fk(l) + fk(r)

def fb(i):
    if i == -1: return 0
    return fh(tree[i][1]) - fh(tree[i][2])

def fr(r):
    l = tree[r][1]
    e = tree[l][2]
    tree[l][2] = r
    tree[r][1] = e
    upd(r)
    upd(l)
    return l

def fl(r):
    l = tree[r][2]
    e = tree[l][1]
    tree[l][1] = r
    tree[r][2] = e
    upd(r)
    upd(l)
    return l

def balance(i):
    b = fb(i)
    if b > 1 and fb(tree[i][1]) >= 0:
        return fr(i)
    if b < -1 and fb(tree[i][2]) <= 0:
        return fl(i)
    if b > 1 and fb(tree[i][1]) < 0:
        tree[i][1] = fl(tree[i][1])
        return fr(i)
    if b < -1 and fb(tree[i][2]) > 0:
        tree[i][2] = fr(tree[i][2])
        return fl(i)
    return i

def add(a, r):
    global tree
    if r == -1:
        i1 = len(tree)
        tree.append([a, -1, -1, 1, 1, 1])
        return i1
    if a < tree[r][0]:
        tree[r][1] = add(a, tree[r][1])
    elif a > tree[r][0]:
        tree[r][2] = add(a, tree[r][2])
    else:
        tree[r][4] += 1
        tree[r][5] += 1
        return r
    upd(r)
    return balance(r)

def find(k, r):
    if r == -1: return None
    lk = fk(tree[r][1])
    if k < lk:
        return find(k, tree[r][1])
    if k < lk + tree[r][4]:
        return tree[r][0]
    return find(k - lk - tree[r][4], tree[r][2])

tree = []
r = -1
n = int(input())

for _ in range(n):
    f, s = input().split()
    if f == "1":
        r = add(s, r)
    else:
        print(find(int(s), r))
⚙️ C++
#include <bits/stdc++.h>
using namespace std;

struct Node {
    int val, left, right, height, cnt, sum;
    Node(int v) : val(v), left(-1), right(-1), height(1), cnt(1), sum(1) {}
};

vector<Node> tree;

int fh(int i) {
    if (i == -1) return 0;
    return tree[i].height;
}

int fk(int i) {
    if (i == -1) return 0;
    return tree[i].sum;
}

void upd(int i) {
    if (i == -1) return;
    int l = tree[i].left;
    int r = tree[i].right;
    tree[i].height = 1 + max(fh(l), fh(r));
    tree[i].sum = tree[i].cnt + fk(l) + fk(r);
}

int fb(int i) {
    if (i == -1) return 0;
    return fh(tree[i].left) - fh(tree[i].right);
}

int fr(int r) {
    int l = tree[r].left;
    int e = tree[l].right;
    tree[l].right = r;
    tree[r].left = e;
    upd(r);
    upd(l);
    return l;
}

int fl(int r) {
    int l = tree[r].right;
    int e = tree[l].left;
    tree[l].left = r;
    tree[r].right = e;
    upd(r);
    upd(l);
    return l;
}

int balance(int i) {
    int b = fb(i);
    if (b > 1 && fb(tree[i].left) >= 0) return fr(i);
    if (b < -1 && fb(tree[i].right) <= 0) return fl(i);
    if (b > 1 && fb(tree[i].left) < 0) {
        tree[i].left = fl(tree[i].left);
        return fr(i);
    }
    if (b < -1 && fb(tree[i].right) > 0) {
        tree[i].right = fr(tree[i].right);
        return fl(i);
    }
    return i;
}

int add(int a, int r) {
    if (r == -1) {
        tree.emplace_back(a);
        return tree.size() - 1;
    }
    if (a < tree[r].val) {
        tree[r].left = add(a, tree[r].left);
    } else if (a > tree[r].val) {
        tree[r].right = add(a, tree[r].right);
    } else {
        tree[r].cnt++;
        tree[r].sum++;
        return r;
    }
    upd(r);
    return balance(r);
}

int find(int k, int r) {
    if (r == -1) return -1;
    int lk = fk(tree[r].left);
    if (k < lk) return find(k, tree[r].left);
    if (k < lk + tree[r].cnt) return tree[r].val;
    return find(k - lk - tree[r].cnt, tree[r].right);
}

int main() {
    ios::sync_with_stdio(0);
    cin.tie(0);
    
    int n;
    cin >> n;
    
    int root = -1;
    
    for (int i = 0; i < n; i++) {
        int f, s;
        cin >> f >> s;
        if (f == 1) {
            root = add(s, root);
        } else {
            cout << find(s, root) << "\n";
        }
    }
    
    return 0;
}

Как работает АВЛ-дерево?

АВЛ-дерево — это самобалансирующееся двоичное дерево поиска (BST). Основная идея: для каждой вершины высоты левого и правого поддеревьев отличаются не более чем на 1 (коэффициент сбалансированности).

Хранимая информация в каждой вершине

Основные функции

Повороты

Балансировка: проверяет баланс и выполняет необходимые повороты (большой левый, большой правый, малые повороты) для восстановления свойства АВЛ-дерева.

Поддержка порядка (k-й элемент)

Алгоритм поддерживает операцию find(k) — поиск k-го по порядку элемента (0-индексация) с использованием заранее вычисленных размеров поддеревьев.

Временная сложность: O(log n) для добавления и O(log n) для поиска k-го элемента.

Пространственная сложность: O(n) для хранения вершин.

Формат ввода и вывода

Ввод:

Пример ввода:

8
1 5
1 3
1 7
1 1
2 0
2 1
2 2
2 3

Вывод:

1
3
5
7

Пояснение:
Добавляем числа: 5, 3, 7, 1
После добавления отсортированный порядок: [1, 3, 5, 7]
0-й элемент = 1, 1-й = 3, 2-й = 5, 3-й = 7

Попробовать АВЛ-дерево онлайн

Введите запросы (1 x — добавить, 2 k — найти k-й элемент):

Результат появится здесь...