Реализация декартова дерева на erlang. Версия на записях.

ctree.erl

-module(ctree).
-export([merge/2, split/2, add/2, add/3, remove/2, flatten/1, random/1, ordered/1, fromlist/1, remove_list/2, depth/1]).

% null / {Key, Param, Left, Right}
-record(cnode, {key, param, left, right}).




%% Объединяет 2 дерева (любой элемент левого меньше любого элемента правого)
% Хотя бы одно из деревьев пустое
merge(null, B) -> B;
merge(A, null) -> A;

% В правое поддерево
merge(#cnode{param=AParam} = A, #cnode{param=BParam} = B) when AParam>BParam ->
  A#cnode{right=merge(A#cnode.right, B)};

% В левое поддерево
merge(A, B) -> 
  B#cnode{left=merge(A, B#cnode.left)}. 


%% Разбивает дерево на 2 поддерева: <SplitKey, >=SplitKey
% Пустое дерево
split(null, _) -> {null, null};

% В правое поддерево
split(#cnode{key=AKey} = A, SplitKey) when AKey<SplitKey ->
  {NL, NR} = split(A#cnode.right, SplitKey), 
  {A#cnode{right=NL}, NR};

% В левое поддерево
split(A, SplitKey) ->
  {NL, NR} = split(A#cnode.left, SplitKey), 
  {NL, A#cnode{left=NR}}.


%% Добавляет ключ Key со случайным параметром в дерево
add(Tree, Key) -> add(Tree, Key, random:uniform()).


%% Добавляет ключ NewKey с параметром NewParam в дерево
% Новое дерево
add(null, NewKey, NewParam) ->
  #cnode{key=NewKey, param=NewParam, left=null, right=null}; 

% Новая вершина
add(#cnode{param=Param} = Tree, NewKey, NewParam) when Param<NewParam -> 
  {NL, NR} = split(Tree, NewKey),
  #cnode{key=NewKey, param=NewParam, left=NL, right=NR};

% В левое поддерево
add(#cnode{key=Key} = Tree, NewKey, NewParam) when Key>NewKey -> 
  Tree#cnode{left=add(Tree#cnode.left, NewKey, NewParam)};

% В правое поддерево
add(Tree, NewKey, NewParam) ->
  Tree#cnode{right=add(Tree#cnode.right, NewKey, NewParam)}.


%% Удаляет ключ RemoveKey из дерева
% Удаляем вершину
remove(#cnode{key=Key, left=Left, right=Right}, RemoveKey) when Key==RemoveKey ->
  merge(Left, Right);

% В левое поддерево
remove(#cnode{key=Key, left=Left} = Tree, RemoveKey) when Key>RemoveKey ->
  Tree#cnode{left=remove(Left, RemoveKey)};

% В правое поддерево
remove(#cnode{right=Right} = Tree, RemoveKey) ->
  Tree#cnode{right=remove(Right, RemoveKey)}.




%%% Мелкие функции для работы с деревьями

flatten(null) -> [];
flatten(#cnode{key=Key, left=Left, right=Right}) -> 
  flatten(Left) ++ [Key | flatten(Right)].

random(0) -> null;
random(Size) -> add(random(Size-1), round(random:uniform()*1000)).

ordered(0) -> null;
ordered(Size) -> add(ordered(Size-1), Size).

fromlist([]) -> null;
fromlist([Head | Tail]) -> add(fromlist(Tail), Head).

remove_list(Tree, []) -> Tree;
remove_list(Tree, [Head | Tail]) -> remove_list(remove(Tree, Head), Tail).
 
depth(null) -> 0;
depth(#cnode{left=Left, right=Right}) -> 
  D1 = depth(Left),
  D2 = depth(Right),
  if 
    D1<D2 -> D2 + 1;
    true -> D1 + 1
  end.