Об дерево

[thedeemon]

Понадобились мне тут в хозяйстве деревья, я взял пилу, топор и пошел в гугл, нашел там AA-деревья, про которые говорится, что они по эффективности как красно-черные, но в реализации проще. Это бинарное дерево, у каждой вершины (узла) есть "уровень", ее высота. Узлы с данными есть двух видов: первый вид имеет форму /\, оба поддерева ниже текущей вершины, обозначим его буквой А. Узлы второго имеют форму Г, в них правое поддерево той же высоты, что и текущая вершина, а левое поддерево ниже; обозначим второй вид буквой H. Инвариант для АА-деревьев звучит так:
1. The level of every leaf node is one.
2. The level of every left child is exactly one less than that of its parent.
3. The level of every right child is equal to or one less than that of its parent.
4. The level of every right grandchild is strictly less than that of its grandparent.
5. Every node of level greater than one has two children.

Тут leaf node - это вершина, у которой нет поддеревьев. Четвертое условие гарантирует, что в дереве не будет цепочек из горизонтальных ребер, поэтому, двигаясь от корня к любой вершине, мы будем спускаться на уровень вниз не реже, чем в половине шагов, а значит самый длинный путь в дереве не более чем в два раза длиннее самого короткого, что дает логарифмическое соотношение высоты и общего числа вершин.

Любой нормальный язык позволяет эти инварианты выразить в типах. Вот так, например, это выглядит в Идрисе:

data Tree : Type -> Nat -> Bool -> Type where
  Null  : Tree a 0 False
  ANode : Tree a n b1 -> a -> Tree a n b2 -> Tree a (n+1) False
  HNode : Tree a n b -> a -> Tree a (n+1) False -> Tree a (n+1) True

Алгебраический тип Tree у нас параметризован тремя вещами: типом данных, хранящихся в дереве, натуральным числом, означающим уровень вершины, и булевым значением, отличающим Н-вершины от остальных. Три конструктора:
Null - отсутствующая вершина, аналог null в иных языках. Высота нуль-вершины всегда 0, и это не Н-вершина, так что пометим False'ом.
ANode - A-вершина, имеет три параметра: левое поддерево, хранимое значение и правое поддерево. Поддеревья должны быть одной высоты, и высота получаемой А-вершины на 1 больше, чем у детей. Не Н-вершина, так что тоже пометим ее False.
HNode - H-вершина с тремя аналогичными параметрами, только на этот раз, во-первых, высота правого поддерева такая же, как у получаемой Н-вершины (на 1 больше, чем высота левого поддерева), а во-вторых, правое поддерево не может быть Н-вершиной, поэтому мы требуем False в соответствующем индексе типа правого поддерева.

Это обычный GADT, доступный и в хаскеле с окамлом, полноценных зависимых типов нам тут не требуется. Правда, вот в этом моменте я не на 100% уверен, пусть более опытные товарищи меня поправят.

Проверим, что этот тип действительно описывает нужные инварианты:
1. Листья кодируются вершинами, имеющими Null в качестве поддеревьев, их высота на 1 больше, чем у Null, т.е. ровно 1.
2,3. Выполнены, как уже показано выше.
4. Правым поддеревом Н-вершины не может оказаться другая Н-вершина, и Null там тоже быть не может (высота не позволит), поэтому там может быть только А-вершина, поддеревья которой ниже, а значит правые "внуки" любой вершины ниже текущей.
5. У вершины высоты больше 1 не может быть Null в качестве поддерева, т.к. высота Null - 0, а разница высот у нас либо 0 либо 1. Значит, у всех вершин выше 1 оба поддерева заселены.

Теперь напишем функцию добавления элемента в дерево, а типы не дадут ей нарушить инварианты. Заимплементировать алгоритм из вики и оригинальной статьи про АА-деревья не получится ровно по той же причине: там сперва элемент добавляют, получая дерево с нарушенными инвариантами, а потом их восстанавливают, проводя необходимые перестройки в дереве. В нашем случае так не получится, ни в какой момент у нас не может быть некорректного дерева, типы не позволят. Поэтому будем сразу строить корректное.

Мы собираемся хранить данные в упорядоченном виде, для этого нам надо их сравнивать, поэтому потребуем, чтобы их тип был из класса типов Ord, т.е. умел сравниваться. Получив на вход дерево высоты n, на выходе у нас может быть дерево либо той же высоты n, либо на 1 больше. Окажется ли корнем Н-вершина мы не знаем, поэтому про соответствующий булевый признак можем лишь сказать, что существует такой b1, что служит соответствующим значением у выходной пары. Как я недавно писал, такое утверждение записывается в виде сигма-типа, зависимой пары. В итоге тип нашей функции выглядит так:
insert : Ord a => Tree a n b -> a -> (b1 ** Either (Tree a n b1) (Tree a (n+1) b1))
Читается: для любого типа а, который можно сравнивать, из дерева с элементами этого типа, высотой n и булевым признаком b, а также значения типа а, мы можем получить такую вот пару: существует такой b1, что либо мы имеем дерево высоты n, либо дерево высоты n+1, с элементами типа а и признаком b1.

А вот полный код нашей функции:

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19

insert Null x = (_ ** Right (ANode Null x Null)) insert (ANode l v r) x with (compare x v) | EQ = (_ ** Left (ANode l x r)) | LT with (insert l x) | (_ ** Left lt) = (_ ** Left (ANode lt v r)) | (_ ** Right (ANode a lv b)) = (_ ** Left (HNode a lv (ANode b v r))) | GT with (insert r x) | (_ ** Left rt) = (_ ** Left (ANode l v rt)) | (_ ** Right (ANode a rv b)) = (_ ** Left (HNode l v (ANode a rv b))) insert (HNode l v r) x with (compare x v) | EQ = (_ ** Left (HNode l x r)) | LT with (insert l x) | (_ ** Left lt) = (_ ** Left (HNode lt v r)) | (_ ** Right lt) = (_ ** Right (ANode lt v r)) | GT with (insert r x) | (False ** Left rt) = (_ ** Left (HNode l v rt)) | (True ** Left (HNode a rv b)) = (_ ** Right (ANode (ANode l v a) rv b)) | (_ ** Right (ANode a rv b)) with a | ANode _ _ _ = (_ ** Right (ANode (HNode l v a) rv b))

По строчкам:
1. Если на входе пустое дерево, то создаем одну А-вершину. Ее высота (1) больше, чем у исходной (0), поэтому возвращаем ее с тэгом Right. И дальше везде тоже: если высота дерева увеличилась, возвращаем его с Right, если не увеличилась, то с тэгом Left. Значение для b1 компилятор может вывести сам, поэтому в явном виде его не указываем, пишем просто (_ ** ...).
2. Если на входе А-вершина, то сравниваем значение в ней (v) с тем, что хотим вставить (x), матчимся по результату сравнения.
3. Если они равны, то перестраивать дерево не нужно, возвращаем такую же вершину.
4. Если вставляемое значение меньше, то вставляем его в левое поддерево и матчимся по результату.
5. Если левое поддерево в результате добавления элемента не увеличило свою высоту (о чем нам говорит тэг Left), то и результат не изменит высоту, возвращаем исходную вершину, где левое поддерево заменено только что полученным.
6. Если же левое поддерево выросло, сравнявшись по высоте с текущим, то делаем из него Н-вершину, правым поддеревом которой становится текущее, а ветка, что там была, становится левой веткой текущего. Общая высота не изменилась, Left.
7. Если вставляемое значение х больше v, то вставляем его в правое поддерево, матчимся по результату.
9. Если правое поддерево после вставки выросло и сравнялось высотой с текущим, делаем из текущего Н-вершину, там как раз правое поддерево той же высоты. Общая высота не изменилась.
10-15. Если на входе Н-вершина, то также сравниваем значение в ней со вставляемым, и если они не равны, вставляем в левое или правое поддерево.
16. Если вставляли справа и правое поддерево не выросло, но это не Н-вершина (о чем свидетельствует False в полученной паре), то ее можно оставить правым поддеревом текущей.
17. Если же справа получилась Н-вершина, то поставить ее справа мы не можем, делаем новую А-вершину еще большей высоты. Общая высота выросла, Right.
18-19. Если после вставки правое поддерево выросло, то его новая высота больше высоты текущего, строим новую А-вершину, содержимое текущей поселяется в ее левом поддереве, а левая ветка выросшего дерева оказывается правым соседом текущей, для чего убеждаемся, что это А-нода, ибо только они могут быть правыми соседями (одного уровня).

Самое приятное тут, что благодаря типам закодировать все эти случаи было намного проще, чем кажется: стоит где-нибудь справа от = ошибиться, перепутать HNode и ANode или Left c Right, как компилятор сразу говорит: "так низя, тут высоты не совпадают, а тут вид вершины не подходит". Написать это все с первого раза без ошибок нереально, и если бы компилятор не проверял все эти инварианты, можно было бы три дня провести в отладке. А так, фигак-фигак и в продакшн. Но есть и но:
Типы здесь гарантируют корректную форму дерева, его сбалансированность. Однако в этом коде ничто не гарантирует правильную упорядоченность хранящихся в дереве значений и даже то, что добавляемый элемент действительно присутствует в полученном дереве. Это можно было бы сделать, усложнив типы и добавив доказательств, но кода стало бы заметно больше. Более того, тут не делается проверка на exhaustiveness вариантов, в коде есть ряд инвариантов, открывающихся лишь методом пристального взгляда. Например, каждый раз, когда в результате добавления элемента поддерево вырастает в высоту (тэг Right), это всегда А-вершина. Данный факт отражен в паттерн-матчинге (строки 6, 9, 18), и вариант с Н-вершиной с тэгом Right даже не рассматривается. Просто потому, что по самому коду видно: везде, где мы возвращаем тэг Right, там внутри А-вершина. Другой факт менее очевидный: когда правый сосед в Н-вершине вырастает (с. 18), не только сама новая вершина с Right вида А, но и ее левое поддерево - тоже А, потому после с.19 нет других вариантов матчинга. Этот факт несложно доказать формально, но в данном коде этого не делается.

И поскольку код тут не все доказывает, пришлось его даже позапускать и потестировать. Работает хорошо.
Чтобы лишние подробности типов и, особенно, сложный тип возвращаемого значения не мешались, сделаем простую обертку:

data AATree : Type -> Type where
  AA : Ord a => Tree a n b -> AATree a 

aa_insert : AATree a -> a -> AATree a
aa_insert (AA t) x with (insert t x)
  | (_ ** Left tr) = AA tr
  | (_ ** Right tr) = AA tr

А теперь самое интересное. Тип Tree индексирован целочисленным и булевым значениями, которые активно используются для обеспечения корректности, на фазе написания и компиляции программы, однако в рантайме они не используются вовсе! Мы нигде не матчимся по ним, нигде поведение программы в рантайме от них не зависит. А значит, можно сделать то же, что внутри делает компилятор: убедившись, что они не используются в рантайме, выкинуть их нафиг. В исходном описании АА-деревьев предлагают в вершинах хранить их высоты, и их активно используют при операциях балансировки. Мы же, упростив типы, получим более простую структуру, которая ведет себя точно так же, гарантированно сохраняет сбалансированность, но высоты в узлах не хранит. Запишем такой усушенный вариант на простом ML (в качестве которого возьмем окамл без выкрутасов), построчно переведя с идриса, с минимальным рефакторингом: проверку на равенство значений вынесем в одно место. Вместо Right и Left факт увеличения высоты будем возвращать в виде булевого значения. Получим:

type node_kind = A | H 
type 'a tree = Null | Node of node_kind * 'a tree * 'a * 'a tree  
 
let rec insert t x = match t with 
  | Null -> true, Node(A, Null, x, Null) 
  | Node(k,l,v,r) -> 
      if x = v then false, Node(k,l,x,r) else 
      match k with 
      | A ->  
        if x < v then  
          (match insert l x with 
          | false, lt -> false, Node(A,lt,v,r) 
          | true, Node(A, a,lv,b) -> false, Node(H, a, lv, Node(A,b,v,r)) ) 
        else 
          (match insert r x with 
          | false, rt -> false, Node(A,l,v,rt) 
          | true, Node(A, a,rv,b) -> false, Node(H,l,v, Node(A,a,rv,b)) ) 
      | H ->  
        if x < v then  
          (match insert l x with 
          | false, lt -> false, Node(H, lt, v, r) 
          | true, lt -> true, Node(A,lt,v,r) ) 
        else 
          (match insert r x with 
          | false, (Node(A,_,_,_) as rt) -> false, Node(H,l,v,rt) 
          | false, Node(H,a,rv,b) -> true, Node(A, Node(A,l,v,a),rv,b) 
          | true, Node(A, a,rv,b) -> true, Node(A, Node(H,l,v,a),rv,b) ) 
 
let aa_insert t x = insert t x |> snd 

Тут компилятор выдаст несколько ворнингов, что не все варианты матчинга рассмотрены, но мы уже убедились, что нерассмотренные варианты и не могут встречаться. Выглядит несложно, но если бы я сразу писал на окамле, то наверняка понаделал бы кучу ошибок и ненужных случаев, потом долго-долго бы отлаживался. Этот код я тоже позапускал и потестил, все работает гладко. Причем в этом варианте становится уже совсем очевидно, как то же самое записать на Си-подобных языках. Тут перевод сделан вручную, следующим этапом будет исходную функцию записать на некоем DSL внутри идриса, чтобы проверить все типы и инварианты, а потом программа бы из AST этого DSL'я сгенерила бы ML или Си-образное что-то.

Чего я пока не делал - это удаления элементов. Оно чуть сложнее, но не сильно. Возможно, там таки понадобится хранить высоты в узлах, пока не знаю.

Comments

[ichernetsky.livejournal.com]

Red-Black Trees in a Functional Settings by Chris Okasaki не подошли?

[thedeemon.livejournal.com]

В конце структура получилась такая же, а код другой: у меня сразу строится корректное дерево, а у Окасаки сперва добавление, потом балансировка. Второй подход проще, но обеспечить корректность там заметно сложнее. Я до этого видел код красно-черных на Омеге, где тоже типами обеспечивают корректность, он был сильно длинный, потому и не стал их копать.