Erlang implementation of an AVL tree

avltree.rb

-module(avltree).
-export([empty_tree/0, insert/3, lookup/2, write_tree/1]).

lookup(Key, { nil, nil, 0, nil, nil }) ->
  not_found;
  
lookup(Key, { Key, Value, _, _, _ }) ->
  { found, Value };
  
lookup(Key, { Key1, _, _, Smaller, _ }) when Key < Key1 ->
  lookup(Key, Smaller);

lookup(Key, { Key1, _, _, Bigger }) when Key > Key1 ->
  lookup(Key, Bigger).
  
insert(Key, Value, { nil, nil, 0, nil, nil }) ->
  E = empty_tree(),
  { Key, Value, 1, E, E };
  
insert(Key, Value, { K2, V2, H2, S2, B2 }) when Key == K2 ->
  { Key, Value, H2, S2, B2 };
  
insert(Key, Value, { K2, V2, H2, S2, B2 }) when Key < K2 ->
  { K4, V4, _, S4, B4 } = insert(Key, Value, S2),
  combine(S4, K4, V4, B4, K2, V2, B2);
  
insert(Key, Value, { K2, V2, H2, S2, B2 }) when Key > K2 ->
  { K4, V4, _, S4, B4 } = insert(Key, Value, B2),
  combine(S2, K2, V2, S4, K4, V4, B4).
  
empty_tree() ->
  { nil, nil, 0, nil, nil }.
  
combine({ K1, V1, H1, S1, B1 }, AK, AV,
        { K2, V2, H2, S2, B2 }, BK, BV,
        { K3, V3, H3, S3, B3 }) when H2 > H1, H2 > H3 ->
          { K2, V2, H1 + 2,
            { AK, AV, H1 + 1, { K1, V1, H1, S1, B1 }, S2 },
            { BK, BV, H3 + 1, B2, { K3, V3, H3, S3, B3 } }
          };
          
combine({ K1, V1, H1, S1, B1 }, AK, AV,
        { K2, V2, H2, S2, B2 }, BK, BV,
        { K3, V3, H3, S3, B3 }) when H1 >= H2, H1 >= H3 ->
          HB = max_add_1(H2, H3),
          HA = max_add_1(H1, HB),
          { AK, AV, HA,
            { K1, V1, H1, S1, B1 },
            { BK, BV, HB, { K2, V2, H2, S2, B2 }, { K3, V3, H3, S3, B3 } }
          };
          
combine({ K1, V1, H1, S1, B1 }, AK, AV,
        { K2, V2, H2, S2, B2 }, BK, BV,
        { K3, V3, H3, S3, B3 }) when H3 >= H1, H3 >= H2 ->
          HA = max_add_1(H1, H2),
          HB = max_add_1(HA, H3),
          { BK, BV, HB,
            { AK, AV, HA, { K1, V1, H1, S1, B1 }, { K2, V2, H2, S2, B2 } },
            { K3, V3, H3, S3, B3 }
          }.
          
max_add_1(X, Y) when Y >= X ->
  Y + 1;

max_add_1(X, Y) when X > Y ->
  X + 1.

write_tree(T) ->
  write_tree(0, T).
  
write_tree(D, { nil, nil, 0, nil, nil }) -> 
  tab(D), 
  io:format('nil', []); 
  write_tree(D, { Key, Value, _, Smaller, Bigger }) -> 
  D1 = D + 4, 
  write_tree(D1, Bigger), 
  io:format('~n', []), 
  tab(D), 
  io:format('~w ===> ~w~n', [ Key, Value ]), 
  write_tree(D1, Smaller). 
  
tab(N) ->
  io:format("~s", [lists:duplicate(N, $ )]).