Trees in Erlang

Like every man and his dog I've been playing around with erlang (and a smattering of Haskell for good measure :). I thought I'd try to see how I might write a binary tree implementation Erlang.

I must confess it turned out much less elegant than I thought it might - maybe some erlang or other functional gurus can point out a better algorithm.

Now before you all comment, yes I know most of the below isn't tail recursive - but performance isn't really the issue here, and a binary tree implementation in erlang is going to be inefficient anyway if you're making frequent updates, since every recursion will involve copying some of the structure (maybe it optimises that away ... dunno enough about its internals to comment).

-module(tree).
-export([find/2, tree_of_squares/1, insert/3]).

-record(node, {key, val, left = nil, right = nil, height = 1}).

%% find/1

find(SearchKey, Node) when SearchKey == Node#node.key ->
  {ok, Node#node.val};

find(_, nil) ->
       not_found;

find(SearchKey, Node) when SearchKey < Node#node.key ->
  find(SearchKey, Node#node.left);

find(SearchKey, Node) ->
  find(SearchKey, Node#node.right).


%% tree_of_squares/1 : make a tree where key = val and contains
%% squares of all numbers from 1 to Max


tree_of_squares(Max) ->
  tree_of_squares(Max, nil).


tree_of_squares(0, Tree) ->
  Tree;

tree_of_squares(Max, Tree) ->
  insert(Max * Max, Max * Max, tree_of_squares(Max - 1, Tree)).

%% insert/3 : return a new tree made from inserting (or updating)
%% a node of Key, Val into existing tree or new tree (nil)

insert(Key, Val, nil) ->
       #node{ key = Key, val = Val };

insert(Key, Val, Node) when Key == Node#node.key ->
       Node#node{ val = Val };

insert(Key, Val, Node) when Key < Node#node.key ->
       balance(Node#node{ left = insert(Key, Val, Node#node.left)});

insert(Key, Val, Node) ->
       balance(Node#node{ right = insert(Key, Val, Node#node.right)}).

%% height/1 : return the height of the tree

height(nil) ->
       0;
height(Node) when Node#node.left == nil, Node#node.right == nil ->
       1;

height(Node) ->
       height(1, Node#node.left, Node#node.right).

height(Acc, Left, Right)
 when Left == nil;
      Right#node.height > Left#node.height ->
       Acc + Right#node.height;

height(Acc, Left, _) ->
       Acc + Left#node.height.


%% balance/1 : return a new tree which is the balanced version of Tree
%% note also that balance is trusted to always return the root height set
%% correctly

balance(nil) ->
  nil;

balance(Node) when Node#node.left == nil, Node#node.right == nil ->
       Node#node{ height = height(Node)};

balance(#node{ left = L, right = R} = Node) when L#node.height ==
R#node.height ->
       Node#node{ height = height(1, L, R)};

balance(#node{ left = L, right = R} = Node)
 when R == nil ;
      L#node.height + 1 > R#node.height ->
       swing_right(Node);


balance(#node{ left = L, right = R} = Node)
 when L == nil ;
      R#node.height + 1 > L#node.height ->
  swing_left(Node).

swing_left(#node{right = #node{ left = RL, right = RR } = Right} = Node)

 when RR == nil ; RR == nil, RL == nil ;
      RL#node.height > RR#node.height ->

       move_left(Node#node{right = move_right(Right)});

swing_left(Tree) ->
       move_left(Tree).

swing_right(#node{left = #node{ left = LL, right = LR} = Left} = Node)

 when LL == nil ; LL == nil, LR == nil ;
      LR#node.height > LL#node.height ->

       move_right(Node#node{ left = move_left(Left)});

swing_right(Tree) ->
       move_right(Tree).


move_left(Node)
 when Node == nil ;
      Node#node.left == nil, Node#node.right == nil ->
       Node;

move_left(#node{ right = #node{ left = RL } = Right} = Node) ->

       OldRoot = Node#node{ right = RL },
       NewRoot = Right#node{ left = OldRoot#node{ height = height(OldRoot)}},
       NewRoot#node{ height = height(NewRoot) }.


move_right(Node)
 when Node == nil ;
      Node#node.left == nil, Node#node.right == nil ->
       Node;

move_right(#node{ left = #node{ right = LR } = Left } = Node) ->

       OldRoot = Node#node{ left = LR},
       NewRoot = Left#node{ right = OldRoot#node{ height = height(OldRoot)}},
       NewRoot#node{ height = height(NewRoot) }.

tree_of_squares is just to easily create an interesting tree to play with, for instance:

1> c(tree).
{ok,tree}
2> T = tree:tree_of_squares(7).
{node,25,
      25,
      {node,9,
            9,
            {node,4,4,{node,1,1,nil,nil,1},nil,2},
            {node,16,16,nil,nil,1},
            3},
      {node,36,36,nil,{node,49,49,nil,nil,1},2},
      4}

Then you can try to find a node:

3> tree:find(1, T).
{ok,1}
4> tree:find(2, T).
not_found

Not very interesting when the key = val - let's insert and update some entries:

5> T2 = tree:insert(5, bar, tree:insert(4, foo, T)).
{node,25,
      25,
      {node,9,
            9,
            {node,4,foo,{node,1,1,nil,nil,1},{node,5,bar,nil,nil,1},2},
            {node,16,16,nil,nil,1},
            3},
      {node,36,36,nil,{node,49,49,nil,nil,1},2},
      4}
6> tree:find(5, T2).
{ok,bar}

Fun huh!

An implementation of red-black trees is left as an exercise for the reader ;)