From 257c2ad71cd4920162ad5df5634c4e81b172b321 Mon Sep 17 00:00:00 2001 From: Simon Cruanes Date: Mon, 7 Sep 2015 22:25:06 +0200 Subject: [PATCH] improve a bit the balancing --- src/data/CCWBTree.ml | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/src/data/CCWBTree.ml b/src/data/CCWBTree.ml index 4237df32..89e3740e 100644 --- a/src/data/CCWBTree.ml +++ b/src/data/CCWBTree.ml @@ -221,6 +221,14 @@ module MakeFull(K : KEY) : S with type key = K.t = struct M.cardinal m = List.length l) *) + (* extract min binding of the tree *) + let rec extract_min_ m = match m with + | E -> assert false + | N (k, v, E, r, _) -> k, v, r + | N (k, v, l, r, _) -> + let k', v', l' = extract_min_ l in + k', v', balance_l k v l' r + (* extract max binding of the tree *) let rec extract_max_ m = match m with | E -> assert false @@ -239,10 +247,16 @@ module MakeFull(K : KEY) : S with type key = K.t = struct | E, o | o, E -> o | _, _ -> - (* remove max element of [l] and put it at the root, - then rebalance towards the left if needed *) - let k', v', l' = extract_max_ l in - balance_l k' v' l' r + if weight l > weight r + then + (* remove max element of [l] and put it at the root, + then rebalance towards the left if needed *) + let k', v', l' = extract_max_ l in + balance_l k' v' l' r + else + (* remove min element of [r] and rebalance *) + let k', v', r' = extract_min_ r in + balance_r k' v' l r' end | n when n<0 -> balance_l k' v' (remove k l) r | _ -> balance_r k' v' l (remove k r)