498 lines
15 KiB
Gleam
498 lines
15 KiB
Gleam
import gleam/option.{type Option, None, Some}
|
||
import gleam/order.{type Order}
|
||
|
||
// Цвета узлов красно-черного дерева
|
||
// Red - красный узел, Black - чёрный узел
|
||
pub type Color {
|
||
Red
|
||
Black
|
||
}
|
||
|
||
// Ленивое значение для отложенных вычислений
|
||
// Позволяет откладывать создание поддеревьев до момента их использования
|
||
// Не знаю правильно это или нет на самом деле, в gleam вроде lazy нету встроенного
|
||
pub type Lazy(a) {
|
||
Thunk(fn() -> a)
|
||
// Отложенное вычисление через функцию
|
||
Value(a)
|
||
// Уже вычисленное значение
|
||
}
|
||
|
||
// Красно-черное дерево с ленивыми вычислениями
|
||
// Полиморфная структура данных с ключами типа k и значениями типа v
|
||
// Поддерживает инварианты красно-чёрного дерева для гарантии балансировки
|
||
pub type RBTree(k, v) {
|
||
Empty
|
||
// Пустое дерево
|
||
Node(
|
||
// Узел дерева
|
||
color: Color,
|
||
// Цвет узла (красный или чёрный)
|
||
key: k,
|
||
// Ключ для поиска и сортировки
|
||
value: v,
|
||
// Значение, связанное с ключом
|
||
left: Lazy(RBTree(k, v)),
|
||
// Левое поддерево (ленивое)
|
||
right: Lazy(RBTree(k, v)),
|
||
// Правое поддерево (ленивое)
|
||
)
|
||
}
|
||
|
||
// Форсирует вычисление ленивого значения
|
||
// Если значение уже вычислено, возвращает его, иначе выполняет функцию
|
||
fn force(lazy: Lazy(a)) -> a {
|
||
case lazy {
|
||
Value(val) -> val
|
||
Thunk(f) -> f()
|
||
}
|
||
}
|
||
|
||
// Создаёт ленивое значение из функции
|
||
// Отложенное вычисление будет выполнено при первом обращении
|
||
fn delay(f: fn() -> a) -> Lazy(a) {
|
||
Thunk(f)
|
||
}
|
||
|
||
// Создаёт пустое дерево - нейтральный элемент моноида
|
||
// Время выполнения: O(1)
|
||
pub fn empty() -> RBTree(k, v) {
|
||
Empty
|
||
}
|
||
|
||
// Проверяет, является ли дерево пустым
|
||
// Время выполнения: O(1)
|
||
pub fn is_empty(tree: RBTree(k, v)) -> Bool {
|
||
case tree {
|
||
Empty -> True
|
||
_ -> False
|
||
}
|
||
}
|
||
|
||
// Создаёт узел с чёрным цветом
|
||
fn make_black(tree: RBTree(k, v)) -> RBTree(k, v) {
|
||
case tree {
|
||
Node(_, key, value, left, right) -> Node(Black, key, value, left, right)
|
||
Empty -> Empty
|
||
}
|
||
}
|
||
|
||
// Балансировка RB tree
|
||
fn balance(
|
||
color: Color,
|
||
key: k,
|
||
value: v,
|
||
left: Lazy(RBTree(k, v)),
|
||
right: Lazy(RBTree(k, v)),
|
||
) -> RBTree(k, v) {
|
||
let left_tree = force(left)
|
||
let right_tree = force(right)
|
||
|
||
case color, left_tree, right_tree {
|
||
// Случай 1: красный левый parent с красными child
|
||
Black, Node(Red, lk, lv, ll, lr), r ->
|
||
case force(ll) {
|
||
Node(Red, llk, llv, lll, llr) ->
|
||
Node(
|
||
Red,
|
||
lk,
|
||
lv,
|
||
delay(fn() { Node(Black, llk, llv, lll, llr) }),
|
||
delay(fn() { Node(Black, key, value, lr, Value(r)) }),
|
||
)
|
||
_ ->
|
||
case force(lr) {
|
||
Node(Red, lrk, lrv, lrl, lrr) ->
|
||
Node(
|
||
Red,
|
||
lrk,
|
||
lrv,
|
||
delay(fn() { Node(Black, lk, lv, ll, lrl) }),
|
||
delay(fn() { Node(Black, key, value, lrr, Value(r)) }),
|
||
)
|
||
_ -> Node(color, key, value, left, right)
|
||
}
|
||
}
|
||
|
||
// Случай 3: красный правый parent с красным левым child
|
||
Black, l, Node(Red, rk, rv, rl, rr) ->
|
||
case force(rl) {
|
||
Node(Red, rlk, rlv, rll, rlr) ->
|
||
Node(
|
||
Red,
|
||
rlk,
|
||
rlv,
|
||
delay(fn() { Node(Black, key, value, Value(l), rll) }),
|
||
delay(fn() { Node(Black, rk, rv, rlr, rr) }),
|
||
)
|
||
_ ->
|
||
case force(rr) {
|
||
Node(Red, rrk, rrv, rrl, rrr) ->
|
||
Node(
|
||
Red,
|
||
rk,
|
||
rv,
|
||
delay(fn() { Node(Black, key, value, Value(l), rl) }),
|
||
delay(fn() { Node(Black, rrk, rrv, rrl, rrr) }),
|
||
)
|
||
_ -> Node(color, key, value, left, right)
|
||
}
|
||
}
|
||
|
||
// Базовый случай - балансировка не нужна
|
||
_, _, _ -> Node(color, key, value, left, right)
|
||
}
|
||
}
|
||
|
||
// Вставляет элемент в дерево с сохранением инвариантов красно-чёрного дерева
|
||
// Время выполнения: O(log n)
|
||
// compare - функция сравнения ключей, должна возвращать order.Lt, order.Eq или order.Gt
|
||
pub fn insert(
|
||
tree: RBTree(k, v),
|
||
key: k,
|
||
value: v,
|
||
compare: fn(k, k) -> Order,
|
||
) -> RBTree(k, v) {
|
||
let result = insert_helper(tree, key, value, compare)
|
||
make_black(result)
|
||
// гарантия что корень чёрный
|
||
}
|
||
|
||
fn insert_helper(
|
||
tree: RBTree(k, v),
|
||
key: k,
|
||
value: v,
|
||
compare: fn(k, k) -> Order,
|
||
) -> RBTree(k, v) {
|
||
case tree {
|
||
Empty -> Node(Red, key, value, Value(Empty), Value(Empty))
|
||
Node(color, k, v, left, right) ->
|
||
case compare(key, k) {
|
||
order.Lt ->
|
||
balance(
|
||
color,
|
||
k,
|
||
v,
|
||
delay(fn() { insert_helper(force(left), key, value, compare) }),
|
||
right,
|
||
)
|
||
order.Gt ->
|
||
balance(
|
||
color,
|
||
k,
|
||
v,
|
||
left,
|
||
delay(fn() { insert_helper(force(right), key, value, compare) }),
|
||
)
|
||
order.Eq -> Node(color, key, value, left, right)
|
||
}
|
||
}
|
||
}
|
||
|
||
// Поиск элемента в дереве
|
||
pub fn lookup(
|
||
tree: RBTree(k, v),
|
||
key: k,
|
||
compare: fn(k, k) -> Order,
|
||
) -> Option(v) {
|
||
case tree {
|
||
Empty -> None
|
||
Node(_, k, v, left, right) ->
|
||
case compare(key, k) {
|
||
order.Lt -> lookup(force(left), key, compare)
|
||
order.Gt -> lookup(force(right), key, compare)
|
||
order.Eq -> Some(v)
|
||
}
|
||
}
|
||
}
|
||
|
||
// Находит минимальный элемент в дереве
|
||
fn find_min(tree: RBTree(k, v)) -> Option(#(k, v)) {
|
||
case tree {
|
||
Empty -> None
|
||
Node(_, key, value, left, _) ->
|
||
case force(left) {
|
||
Empty -> Some(#(key, value))
|
||
_ -> find_min(force(left))
|
||
}
|
||
}
|
||
}
|
||
|
||
// Удаляет минимальный элемент из дерева
|
||
fn delete_min(tree: RBTree(k, v)) -> RBTree(k, v) {
|
||
case tree {
|
||
Empty -> Empty
|
||
Node(color, key, value, left, right) ->
|
||
case force(left) {
|
||
Empty -> force(right)
|
||
_ ->
|
||
Node(
|
||
color,
|
||
key,
|
||
value,
|
||
delay(fn() { delete_min(force(left)) }),
|
||
right,
|
||
)
|
||
}
|
||
}
|
||
}
|
||
|
||
pub fn delete(
|
||
tree: RBTree(k, v),
|
||
key: k,
|
||
compare: fn(k, k) -> Order,
|
||
) -> RBTree(k, v) {
|
||
case tree {
|
||
Empty -> Empty
|
||
Node(_, k, v, left, right) ->
|
||
case compare(key, k) {
|
||
order.Lt ->
|
||
Node(
|
||
Black,
|
||
k,
|
||
v,
|
||
delay(fn() { delete(force(left), key, compare) }),
|
||
right,
|
||
)
|
||
order.Gt ->
|
||
Node(
|
||
Black,
|
||
k,
|
||
v,
|
||
left,
|
||
delay(fn() { delete(force(right), key, compare) }),
|
||
)
|
||
order.Eq ->
|
||
case force(left), force(right) {
|
||
Empty, Empty -> Empty
|
||
_, Empty -> force(left)
|
||
Empty, _ -> force(right)
|
||
_, _ ->
|
||
case find_min(force(right)) {
|
||
None -> force(left)
|
||
Some(#(min_key, min_value)) ->
|
||
Node(
|
||
Black,
|
||
min_key,
|
||
min_value,
|
||
left,
|
||
delay(fn() { delete_min(force(right)) }),
|
||
)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// Фильтрация элементов дерева
|
||
pub fn filter(
|
||
tree: RBTree(k, v),
|
||
predicate: fn(k, v) -> Bool,
|
||
compare: fn(k, k) -> Order,
|
||
) -> RBTree(k, v) {
|
||
fold_left(tree, empty(), fn(acc, key, value) {
|
||
case predicate(key, value) {
|
||
True -> insert(acc, key, value, compare)
|
||
False -> acc
|
||
}
|
||
})
|
||
}
|
||
|
||
// Отображение значений в дереве
|
||
pub fn map(tree: RBTree(k, v), f: fn(v) -> w) -> RBTree(k, w) {
|
||
case tree {
|
||
Empty -> Empty
|
||
Node(color, key, value, left, right) ->
|
||
Node(
|
||
color,
|
||
key,
|
||
f(value),
|
||
delay(fn() { map(force(left), f) }),
|
||
delay(fn() { map(force(right), f) }),
|
||
)
|
||
}
|
||
}
|
||
|
||
// Левая свёртка дерева
|
||
pub fn fold_left(tree: RBTree(k, v), acc: a, f: fn(a, k, v) -> a) -> a {
|
||
case tree {
|
||
Empty -> acc
|
||
Node(_, key, value, left, right) -> {
|
||
let left_acc = fold_left(force(left), acc, f)
|
||
let current_acc = f(left_acc, key, value)
|
||
fold_left(force(right), current_acc, f)
|
||
}
|
||
}
|
||
}
|
||
|
||
// Правая свёртка дерева
|
||
pub fn fold_right(tree: RBTree(k, v), acc: a, f: fn(k, v, a) -> a) -> a {
|
||
case tree {
|
||
Empty -> acc
|
||
Node(_, key, value, left, right) -> {
|
||
let right_acc = fold_right(force(right), acc, f)
|
||
let current_acc = f(key, value, right_acc)
|
||
fold_right(force(left), current_acc, f)
|
||
}
|
||
}
|
||
}
|
||
|
||
// Объединение двух деревьев (операция моноида)
|
||
pub fn concat(
|
||
tree1: RBTree(k, v),
|
||
tree2: RBTree(k, v),
|
||
compare: fn(k, k) -> Order,
|
||
) -> RBTree(k, v) {
|
||
fold_left(tree2, tree1, fn(acc, key, value) {
|
||
insert(acc, key, value, compare)
|
||
})
|
||
}
|
||
|
||
// Нейтральный элемент моноида (пустое дерево)
|
||
pub fn mempty() -> RBTree(k, v) {
|
||
empty()
|
||
}
|
||
|
||
// Размер дерева
|
||
pub fn size(tree: RBTree(k, v)) -> Int {
|
||
fold_left(tree, 0, fn(acc, _, _) { acc + 1 })
|
||
}
|
||
|
||
// Преобразование дерева в список пар
|
||
pub fn to_list(tree: RBTree(k, v)) -> List(#(k, v)) {
|
||
fold_right(tree, [], fn(key, value, acc) { [#(key, value), ..acc] })
|
||
}
|
||
|
||
// Создание дерева из списка пар
|
||
pub fn from_list(
|
||
list: List(#(k, v)),
|
||
compare: fn(k, k) -> Order,
|
||
) -> RBTree(k, v) {
|
||
list
|
||
|> list_fold_left(empty(), fn(acc, pair) {
|
||
let #(key, value) = pair
|
||
insert(acc, key, value, compare)
|
||
})
|
||
}
|
||
|
||
// Вспомогательная функция для свёртки списка
|
||
fn list_fold_left(list: List(a), acc: b, f: fn(b, a) -> b) -> b {
|
||
case list {
|
||
[] -> acc
|
||
[head, ..tail] -> list_fold_left(tail, f(acc, head), f)
|
||
}
|
||
}
|
||
|
||
// Эффективная проверка равенства двух деревьев
|
||
// Сравнивает деревья структурно без преобразования в списки
|
||
// Время выполнения: O(min(n, m)) где n, m - размеры деревьев
|
||
pub fn equal(
|
||
tree1: RBTree(k, v),
|
||
tree2: RBTree(k, v),
|
||
key_compare: fn(k, k) -> Bool,
|
||
value_compare: fn(v, v) -> Bool,
|
||
) -> Bool {
|
||
equal_helper(tree1, tree2, key_compare, value_compare)
|
||
}
|
||
|
||
// Вспомогательная функция для структурного сравнения деревьев
|
||
// Использует замыкание при первом несовпадении
|
||
fn equal_helper(
|
||
tree1: RBTree(k, v),
|
||
tree2: RBTree(k, v),
|
||
key_compare: fn(k, k) -> Bool,
|
||
value_compare: fn(v, v) -> Bool,
|
||
) -> Bool {
|
||
case tree1, tree2 {
|
||
// Оба дерева пустые - равны
|
||
Empty, Empty -> True
|
||
|
||
Empty, _ -> False
|
||
_, Empty -> False
|
||
|
||
Node(color1, key1, value1, left1, right1),
|
||
Node(color2, key2, value2, left2, right2)
|
||
-> {
|
||
color1 == color2
|
||
// Ключи должны быть равны
|
||
&& key_compare(key1, key2)
|
||
// Значения должны быть равны
|
||
&& value_compare(value1, value2)
|
||
// Рекурсивно проверяем левые поддеревья (ленивые)
|
||
&& equal_helper(force(left1), force(left2), key_compare, value_compare)
|
||
// Рекурсивно проверяем правые поддеревья (ленивые)
|
||
&& equal_helper(force(right1), force(right2), key_compare, value_compare)
|
||
}
|
||
}
|
||
}
|
||
|
||
// Проверка, содержится ли ключ в дереве
|
||
pub fn contains(tree: RBTree(k, v), key: k, compare: fn(k, k) -> Order) -> Bool {
|
||
case lookup(tree, key, compare) {
|
||
Some(_) -> True
|
||
None -> False
|
||
}
|
||
}
|
||
|
||
// Получение всех ключей дерева
|
||
pub fn keys(tree: RBTree(k, v)) -> List(k) {
|
||
fold_right(tree, [], fn(key, _, acc) { [key, ..acc] })
|
||
}
|
||
|
||
// Получение всех значений дерева
|
||
pub fn values(tree: RBTree(k, v)) -> List(v) {
|
||
fold_right(tree, [], fn(_, value, acc) { [value, ..acc] })
|
||
}
|
||
|
||
// Семантическое равенство деревьев - сравнивает содержимое независимо от структуры
|
||
// Два дерева семантически равны, если содержат одинаковые пары ключ-значение
|
||
// Время выполнения: O(n log n) из-за необходимости итерации
|
||
pub fn semantic_equal(
|
||
tree1: RBTree(k, v),
|
||
tree2: RBTree(k, v),
|
||
key_order: fn(k, k) -> Order,
|
||
value_equal: fn(v, v) -> Bool,
|
||
) -> Bool {
|
||
// Быстрая проверка размеров
|
||
case size(tree1) == size(tree2) {
|
||
False -> False
|
||
True -> {
|
||
// Проверяем, что все элементы из tree1 есть в tree2 с теми же значениями
|
||
fold_left(tree1, True, fn(acc, key, value) {
|
||
acc
|
||
&& case lookup(tree2, key, key_order) {
|
||
Some(other_value) -> value_equal(value, other_value)
|
||
None -> False
|
||
}
|
||
})
|
||
}
|
||
}
|
||
}
|
||
|
||
// Сравнение с игнорированием цветов узлов (только структура и данные)
|
||
// Полезно когда важна только логическая структура дерева поиска
|
||
pub fn structure_equal(
|
||
tree1: RBTree(k, v),
|
||
tree2: RBTree(k, v),
|
||
key_compare: fn(k, k) -> Bool,
|
||
value_compare: fn(v, v) -> Bool,
|
||
) -> Bool {
|
||
case tree1, tree2 {
|
||
Empty, Empty -> True
|
||
Empty, _ -> False
|
||
_, Empty -> False
|
||
Node(_, key1, value1, left1, right1), Node(_, key2, value2, left2, right2) -> {
|
||
// Игнорируем цвет, сравниваем только ключи, значения и структуру
|
||
key_compare(key1, key2)
|
||
&& value_compare(value1, value2)
|
||
&& structure_equal(force(left1), force(left2), key_compare, value_compare)
|
||
&& structure_equal(
|
||
force(right1),
|
||
force(right2),
|
||
key_compare,
|
||
value_compare,
|
||
)
|
||
}
|
||
}
|
||
}
|