erlang的gb_trees源码解析

Erlang的gb_trees给出了平衡二叉树的实现。

 

empty()函数给出了一个得到一个空的二叉树的途径。

empty() ->
    {0, nil}.

空树作为一个含有两个成员的元组，第一个0则代表树中的元素个数，第二个位置则用来存放根节点。

 

insert()函数则给出了给二叉树添加节点的途径。

insert(Key, Val, {S, T}) when is_integer(S) ->
    S1 = S+1,
    {S1, insert_1(Key, Val, T, ?pow(S1, ?p))}.

insert_1(Key, Value, {Key1, V, Smaller, Bigger}, S) when Key < Key1 -> 
    case insert_1(Key, Value, Smaller, ?div2(S)) of
 {T1, H1, S1} ->
     T = {Key1, V, T1, Bigger},
     {H2, S2} = count(Bigger),
     H = ?mul2(erlang:max(H1, H2)),
     SS = S1 + S2 + 1,
     P = ?pow(SS, ?p),
     if
  H > P -> 
      balance(T, SS);
  true ->
      {T, H, SS}
     end;
 T1 ->
     {Key1, V, T1, Bigger}
    end;
insert_1(Key, Value, {Key1, V, Smaller, Bigger}, S) when Key > Key1 -> 
    case insert_1(Key, Value, Bigger, ?div2(S)) of
 {T1, H1, S1} ->
     T = {Key1, V, Smaller, T1},
     {H2, S2} = count(Smaller),
     H = ?mul2(erlang:max(H1, H2)),
     SS = S1 + S2 + 1,
     P = ?pow(SS, ?p),
     if
  H > P -> 
      balance(T, SS);
  true ->
      {T, H, SS}
     end;
 T1 ->
     {Key1, V, Smaller, T1}
    end;
insert_1(Key, Value, nil, S) when S =:= 0 ->
    {{Key, Value, nil, nil}, 1, 1};
insert_1(Key, Value, nil, _S) ->
    {Key, Value, nil, nil};
insert_1(Key, _, _, _) ->
    erlang:error({key_exists, Key}).

在使用insert()添加节点的时候，首先给原本元组中表示底下元素的S加一作为S1代表添加进值之后的后节点以及节点下的节点数量。

之后求得S的平方，作为参数通过insert_1()函数添加节点。在第一次添加的情况下爱，第三个参数本身就是nil也就是空节点，那么可以直接将所要添加的节点放到这个位置，构造一个四个元素的元组，前两个就是键值对，而后面则代表左右子节点，当然此时为空。将这个元组返回，所以最后得到的结果为一个两个元素的元组，第一个代表此时所有节点的数量，第二个则是新加入的根节点。

 

如果已经存在节点的情况下通过insert()函数加入新节点，可能较为复杂。

当新加入的key小于根节点的key，那么得到当前节点的右节点继续递归执行insert_1()函数进行添加，如果是空的，那么直接加入。

同理，当新加入的key大于当前节点的key，那么也是直接当前节点的左节点继续递归执行insert_1()函数进行添加，与之前一样。

 

但是此时的树是不平衡的，只是一个简单的二叉树。那么什么时候开始树的平衡操作呢？

在insert()正式准备添加节点的时候，计算得到了加入新节点的个数的2的平方，这个数字就是用来作为是否进行二叉树平衡的依据。

-define(pow(A, _), A * A).

每次在进行一次新的insert_1()递归的时候，都会将得到的平方右移一位，也就是说新加入节点的2次方的二进制位数代表了递归次数的上限，位移之后为0，也就是说当前的树可能需要平衡操作，这是返回的元组包含了一个新添加进来的节点，以及H1和S1，分别为1。返回得到的结果，先将新得到节点放到原本操作应该在的位置。

count({_, _, nil, nil}) ->
    {1, 1};
count({_, _, Sm, Bi}) ->
    {H1, S1} = count(Sm),
    {H2, S2} = count(Bi),
    {?mul2(erlang:max(H1, H2)), S1 + S2 + 1};
count(nil) ->
    {1, 0}.

count()函数用来计算某个节点下的两个数字，H1的位数则代表当前节点下最大成熟，S1代表当前节点下（包括当前节点）的所有节点数量。

通过conut()函数得到新节点的这两个参数为H2，和S2。将H1与H2最大值再左移一位，H的位数代表当前节点的父节点的最大层数，之后将父节点以及所有节点的个数相加得到SS并求2的次方为P。

如果H大于P，则代表此时的树过于稀疏，需要采用balance()方法将整棵树变成二叉平衡树。

balance(T, S) ->
    balance_list(to_list_1(T), S).

balance_list(L, S) ->
    {T, []} = balance_list_1(L, S),
    T.

balance_list_1(L, S) when S > 1 ->
    Sm = S - 1,
    S2 = Sm div 2,
    S1 = Sm - S2,
    {T1, [{K, V} | L1]} = balance_list_1(L, S1),
    {T2, L2} = balance_list_1(L1, S2),
    T = {K, V, T1, T2},
    {T, L2};
balance_list_1([{Key, Val} | L], 1) ->
    {{Key, Val, nil, nil}, L};
balance_list_1(L, 0) ->
    {nil, L}.

to_list_1(T) -> to_list(T, []).

to_list({Key, Value, Small, Big}, L) ->
    to_list(Small, [{Key, Value} | to_list(Big, L)]);
to_list(nil, L) -> L.

 

传入的参数为当前新节点的父节点和父节点以及其下面所有节点的数量。

首先将当前的树，也就是元组通过to_list_1()转化为列表，结果其实就是将当前二叉树的先序遍历的结果。

之后通过balance_list_1()函数开始转化为二叉平衡树。

参数S为节点下的数量，减一得到Sm，除以2得到S2，S1为Sm减去S2的结果作为当前节点左边的节点个数。而后继续递归balance_list_1()，直到S为1，此时该节点，也就是当前树先序遍历之后的第一个节点就是叶子节点。之后将下一个节点取出，这个节点就是上一个节点的父节点，继续将列表的下一个节点通过balance_list_1()进行递归，由于列表是前序遍历的结果，这样把上述T1，T2分别作为当前节点的左右节点，继续上一层递归，由此平衡二叉树的平衡完成。

 

balance()完成之后，返回的三个成员的元组不止含有当前父节点，还有层数H以及节点数SS便于上层节点的平衡。