Сбалансированное двоичное дерево поиска с поддержкой порядка
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))
#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 (коэффициент сбалансированности).
Алгоритм поддерживает операцию find(k) — поиск k-го по порядку элемента (0-индексация) с использованием заранее вычисленных размеров поддеревьев.
Временная сложность: O(log n) для добавления и O(log n) для поиска k-го элемента.
Пространственная сложность: O(n) для хранения вершин.
Ввод:
n — количество запросов1 x — добавить значение x в дерево2 k — найти k-й по порядку элемент (0-индексация) и вывести его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-й элемент):