Motiejaus dvejetainio paieškos medžio implementacija

btree.erl

-module(btree).
-export([bt/1, lookup/2, ins/2, del/2]).

%% UNBALANCED BINARY TREE
bt(Val, Left, Right) -> { Val, Left, Right }.
bt(Val) -> bt(Val, nil, nil).

lookup(nil, _) -> nil;
lookup({Val, _, _}, Val) -> Val;
lookup({Val, L, _}, Elem) when Val > Elem -> lookup(L, Elem);
lookup({Val, _, R}, Elem) when Val < Elem -> lookup(R, Elem).

ins(nil, Elem) -> bt(Elem);
ins({Val, _, _} = Node, Val) -> Node;
ins({Val, L, R}, Elem) when Elem < Val -> { Val, ins(L, Elem), R };
ins({Val, L, R}, Elem) when Elem > Val -> { Val, L, ins(R, Elem) }.

% Element not found, get out
del(nil, _) -> nil;

% Traversing, Node not found yet.
del({Val, L, R}, Elem) when (Val =/= Elem) ->
    if
        (Elem < Val) -> {Val, del(L, Elem), R};
        (Elem > Val) -> {Val, L, del(R, Elem)}
    end;

% Node found. No children. Just remove node.
del({Val, nil, nil}, Val) -> nil;

% Node found. One child. Replace current element with the child
del({Val, L, R}, Val) when (L =:= nil) or (R =:= nil) ->
    if
        L =:= nil -> R;
        R =:= nil -> L
    end;

% Node found. Two children. Replace current element with largest inode element.
del({Val, L, R}, Val) ->
    {MaxVal, Node} = del_max(L),
    {MaxVal, Node, R}.

% Delete maximum inode element. Returns Node and Max leaf value
% Bottom leaf, has no children. Return the leaf value and empty Node
del_max({MaxVal, nil, nil}) -> {MaxVal, nil};

% Return Node and Max leaf value
del_max({Val, L, nil}) ->
    {MaxVal, Node} = del_max(L),
    {MaxVal, {Val, Node, nil}};

del_max({Val, L, R}) ->
    {MaxVal, Node} = del_max(R),
    {MaxVal, {Val, L, Node}}.

-include_lib("eunit/include/eunit.hrl").

bt_generator_test_() -> [
        { "Root node constructor", ?_assert(bt(5) =:= { 5, nil, nil }) },
        { "Construct root and child nodes",
            ?_assert( bt(5, bt(3), nil) =:= { 5, {3, nil, nil}, nil }) }
    ].

sample() -> bt(5, bt(3, nil, bt(4)), bt(6)).
bt_lookup_test_() -> [
        { "Root node lookup", ?_assert(lookup(sample(), 5) =:= 5 )},
        { "Child node lookup", ?_assert(lookup(sample(), 3) =:= 3 )},
        { "Leaf node lookup", ?_assert(lookup(sample(), 4) =:= 4 )},
        { "N/E lookup", ?_assert(lookup(sample(), 10) =:= nil )},

        { "Lookup in root element", ?_test(lookup(bt(0), 0) =:= 0 )},
        { "N/E in root element", ?_assert(lookup(bt(0), 1) =:= nil )}
    ].
bt_ins_test_() -> [
        { "Insertion to right of root element",
            ?_assert(ins(bt(5), 9) =:= {5, nil, {9, nil, nil}} )},
        { "Insertion to left of root element",
            ?_assert(ins(bt(5), 1) =:= {5, {1, nil, nil}, nil} )}
    ].
bt_del_test_() -> [
        { "Do not modify tree if it does not have that element 1",
            ?_assert(del(bt(5), 1) =:= bt(5) )},
        { "Do not modify tree if it does not have that element 2",
            ?_assert(del(bt(5, bt(1), bt(9)), 0) =:= bt(5, bt(1), bt(9)) )},

        { "Remove childless root node", ?_assert( del(bt(5), 5) =:= nil )},

        { "Remove left childless element from root node",
            ?_assert(del(bt(5, bt(1), bt(9)), 1) =:= bt(5, nil, bt(9)) )},

        { "Remove right childless element from root node",
            ?_assert(del(bt(5, bt(1), bt(9)), 9) =:= bt(5, bt(1), nil) )},

        { "Remove single child element (left child)",
            ?_assert(del(bt(5, bt(1), nil), 1) =:= bt(5) )},

        { "Remove single child element (right child)",
            ?_assert(del(bt(5, nil, bt(9)), 9) =:= bt(5) )},

        { "Remove element with 2 children",
            ?_assert(del(bt(5, bt(1), bt(9)), 5) =:= bt(1, nil, bt(9)) )}
    ].