mirror of
https://github.com/c-cube/ocaml-containers.git
synced 2025-12-07 11:45:31 -05:00
Merge pull request #3 from nbraud/avl
AVL: Change API (why not, it makes for a usable polymorphic map)
This commit is contained in:
Simon Cruanes
commit
b661147c3c
2 changed files with 110 additions and 76 deletions
150
misc/AVL.ml
150
misc/AVL.ml
|
|
@ -28,13 +28,17 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
|||
|
||||
See https://en.wikipedia.org/wiki/AVL_tree *)
|
||||
|
||||
type ('a,'b) t =
|
||||
| Empty
|
||||
| Node of ('a,'b) t * 'a * 'b * ('a,'b) t * int
|
||||
|
||||
type 'a comparator = 'a -> 'a -> int
|
||||
|
||||
let empty = Empty
|
||||
type ('a,'b) tree =
|
||||
| Empty
|
||||
| Node of ('a,'b) tree * 'a * 'b * ('a,'b) tree * int
|
||||
|
||||
type ('a,'b) t =
|
||||
{ cmp: 'a comparator;
|
||||
t: ('a,'b) tree }
|
||||
|
||||
let empty ~cmp = { cmp; t=Empty }
|
||||
|
||||
let _height = function
|
||||
| Empty -> 0
|
||||
|
|
@ -46,7 +50,8 @@ let _balance l r = _height l - _height r
|
|||
let _make l x y r =
|
||||
Node (l, x, y, r, 1 + max (_height l) (_height r))
|
||||
|
||||
let singleton k v = _make empty k v empty
|
||||
let _singleton k v = _make Empty k v Empty
|
||||
let singleton ~cmp k v = { cmp; t = _singleton k v }
|
||||
|
||||
(* balance tree [t] *)
|
||||
let _rebalance t = match t with
|
||||
|
|
@ -93,32 +98,40 @@ let _rebalance t = match t with
|
|||
let _make_balance l k v r =
|
||||
_rebalance (_make l k v r)
|
||||
|
||||
let rec fold f acc t = match t with
|
||||
let rec _fold f acc t = match t with
|
||||
| Empty -> acc
|
||||
| Node (l, x, y, r, _) ->
|
||||
let acc = fold f acc l in
|
||||
let acc = _fold f acc l in
|
||||
let acc = f acc x y in
|
||||
fold f acc r
|
||||
_fold f acc r
|
||||
|
||||
let rec for_all p t = match t with
|
||||
let fold f acc {t; _} = _fold f acc t
|
||||
|
||||
let rec _for_all p t = match t with
|
||||
| Empty -> true
|
||||
| Node (l, x, y, r, _) ->
|
||||
p x y && for_all p l && for_all p r
|
||||
p x y && _for_all p l && _for_all p r
|
||||
|
||||
let rec exists p t = match t with
|
||||
let for_all p {t; _} = _for_all p t
|
||||
|
||||
let rec _exists p t = match t with
|
||||
| Empty -> false
|
||||
| Node (l, x, y, r, _) ->
|
||||
p x y || exists p l || exists p r
|
||||
p x y || _exists p l || _exists p r
|
||||
|
||||
let rec insert ~cmp t k v = match t with
|
||||
| Empty -> _make empty k v empty
|
||||
let exists p {t; _} = _exists p t
|
||||
|
||||
let rec _insert ~cmp t k v = match t with
|
||||
| Empty -> _make Empty k v Empty
|
||||
| Node (l, k1, v1, r, _) ->
|
||||
let c = cmp k k1 in
|
||||
if c < 0
|
||||
then _make_balance (insert ~cmp l k v) k1 v1 r
|
||||
then _make_balance (_insert ~cmp l k v) k1 v1 r
|
||||
else if c = 0
|
||||
then _make l k v r
|
||||
else _make_balance l k1 v1 (insert ~cmp r k v)
|
||||
else _make_balance l k1 v1 (_insert ~cmp r k v)
|
||||
|
||||
let insert {cmp; t} k v = {cmp; t=_insert ~cmp t k v}
|
||||
|
||||
(* remove the maximal value in the given tree (the only which only has a left
|
||||
child), and return its key/value pair *)
|
||||
|
|
@ -132,7 +145,7 @@ let rec _remove_max t = match t with
|
|||
|
||||
exception NoSuchElement
|
||||
|
||||
let remove ~cmp t key =
|
||||
let _remove ~cmp t key =
|
||||
let rec _remove t = match t with
|
||||
| Empty -> raise NoSuchElement
|
||||
| Node (l, k, v, r, _) ->
|
||||
|
|
@ -153,30 +166,34 @@ let remove ~cmp t key =
|
|||
try _remove t
|
||||
with NoSuchElement -> t (* element not found *)
|
||||
|
||||
let update ~cmp t key f = failwith "update: not implemented"
|
||||
let remove {cmp; t} k = {cmp; t=_remove ~cmp t k}
|
||||
|
||||
let rec find_exn ~cmp t key = match t with
|
||||
let _update ~cmp t key f = failwith "update: not implemented"
|
||||
let update {cmp; t} = _update ~cmp t
|
||||
|
||||
let rec _find_exn ~cmp t key = match t with
|
||||
| Empty -> raise Not_found
|
||||
| Node (l, k, v, r, _) ->
|
||||
let c = cmp key k in
|
||||
if c < 0 then find_exn ~cmp l key
|
||||
else if c > 0 then find_exn ~cmp r key
|
||||
if c < 0 then _find_exn ~cmp l key
|
||||
else if c > 0 then _find_exn ~cmp r key
|
||||
else v
|
||||
let find_exn {cmp; t} = _find_exn ~cmp t
|
||||
|
||||
let find ~cmp t key =
|
||||
try Some (find_exn ~cmp t key)
|
||||
let find t key =
|
||||
try Some (find_exn t key)
|
||||
with Not_found -> None
|
||||
|
||||
(* add k,v as strictly maximal element to t. [t] must not contain
|
||||
any key >= k *)
|
||||
let rec _add_max k v t = match t with
|
||||
| Empty -> singleton k v
|
||||
| Empty -> _singleton k v
|
||||
| Node (l, k', v', r, _) ->
|
||||
_make_balance l k' v' (_add_max k v r)
|
||||
and
|
||||
|
||||
(* same for minimal value *)
|
||||
_add_min k v t = match t with
|
||||
| Empty -> singleton k v
|
||||
let rec _add_min k v t = match t with
|
||||
| Empty -> _singleton k v
|
||||
| Node (l, k', v', r, _) ->
|
||||
_make_balance (_add_min k v l) k' v' r
|
||||
|
||||
|
|
@ -204,42 +221,51 @@ let _concat t1 t2 = match t1, t2 with
|
|||
let t1', k, v = _remove_max t1 in
|
||||
_join t1' k v t2
|
||||
|
||||
let rec split ~cmp t key = match t with
|
||||
| Empty -> empty, None, empty
|
||||
let rec _split ~cmp t key = match t with
|
||||
| Empty -> Empty, None, Empty
|
||||
| Node (l, k, v, r, _) ->
|
||||
let c = cmp key k in
|
||||
if c < 0
|
||||
then
|
||||
let ll, result, lr = split ~cmp l key in
|
||||
let ll, result, lr = _split ~cmp l key in
|
||||
ll, result, _join lr k v r
|
||||
else if c > 0
|
||||
then
|
||||
let rl, result, rr = split ~cmp r key in
|
||||
let rl, result, rr = _split ~cmp r key in
|
||||
_join l k v rl, result, rr
|
||||
else
|
||||
l, Some v, r
|
||||
|
||||
let split {cmp; t} k =
|
||||
let (t,b,t') = _split ~cmp t k in
|
||||
{cmp; t}, b, {cmp; t=t'}
|
||||
|
||||
(* if k = Some v, join l k v r, else concat l v *)
|
||||
let _concat_or_join l k result r = match result with
|
||||
| None -> _concat l r
|
||||
| Some v -> _join l k v r
|
||||
|
||||
let rec merge ~cmp f t1 t2 = match t1, t2 with
|
||||
| Empty, Empty -> empty
|
||||
let rec _merge ~cmp f t1 t2 = match t1, t2 with
|
||||
| Empty, Empty -> Empty
|
||||
| Node (l1, k1, v1, r1, h1), _ when h1 >= _height t2 ->
|
||||
let l2, result2, r2 = split ~cmp t2 k1 in
|
||||
let l2, result2, r2 = _split ~cmp t2 k1 in
|
||||
let result = f k1 (Some v1) result2 in
|
||||
let l = merge ~cmp f l1 l2 in
|
||||
let r = merge ~cmp f r1 r2 in
|
||||
let l = _merge ~cmp f l1 l2 in
|
||||
let r = _merge ~cmp f r1 r2 in
|
||||
_concat_or_join l k1 result r
|
||||
| _, Node (l2, k2, v2, r2, _) ->
|
||||
let l1, result1, r1 = split ~cmp t1 k2 in
|
||||
let l1, result1, r1 = _split ~cmp t1 k2 in
|
||||
let result = f k2 result1 (Some v2) in
|
||||
let l = merge ~cmp f l1 l2 in
|
||||
let r = merge ~cmp f r1 r2 in
|
||||
let l = _merge ~cmp f l1 l2 in
|
||||
let r = _merge ~cmp f r1 r2 in
|
||||
_concat_or_join l k2 result r
|
||||
| _, Empty -> assert false (* h1 < heigth h2?? *)
|
||||
|
||||
let merge f {cmp; t} {cmp=cmp'; t=t'} =
|
||||
if(cmp != cmp') then invalid_arg "AVL.merge: trees wit different
|
||||
comparison function";
|
||||
{cmp; t = _merge ~cmp f t t'}
|
||||
|
||||
(* invariant: balanced *)
|
||||
let rec invariant_balance t = match t with
|
||||
| Empty -> true
|
||||
|
|
@ -253,13 +279,13 @@ let rec invariant_search ~cmp t = match t with
|
|||
| Node (l, x, _, r, _) ->
|
||||
invariant_search ~cmp l &&
|
||||
invariant_search ~cmp r &&
|
||||
for_all (fun x' _ -> cmp x' x < 0) l &&
|
||||
for_all (fun x' _ -> cmp x' x > 0) r
|
||||
_for_all (fun x' _ -> cmp x' x < 0) l &&
|
||||
_for_all (fun x' _ -> cmp x' x > 0) r
|
||||
|
||||
let of_list ~cmp l =
|
||||
List.fold_left (fun acc (x,y) -> insert ~cmp acc x y) empty l
|
||||
{cmp; t = List.fold_left (fun acc (x,y) -> _insert ~cmp acc x y) Empty l}
|
||||
|
||||
let to_list t =
|
||||
let to_list {t; _} =
|
||||
let rec aux acc t = match t with
|
||||
| Empty -> acc
|
||||
| Node (l, k, v, r, _) ->
|
||||
|
|
@ -273,15 +299,15 @@ let to_list t =
|
|||
module type ITERATOR = sig
|
||||
type 'a iter
|
||||
|
||||
val after : cmp:'a comparator -> ('a,'b) t -> 'a -> ('a * 'b) iter
|
||||
val before : cmp:'a comparator -> ('a,'b) t -> 'a -> ('a * 'b) iter
|
||||
val iter : ('a,'b) t -> ('a * 'b) iter
|
||||
val add : cmp:'a comparator -> ('a,'b) t -> ('a * 'b) iter -> ('a,'b) t
|
||||
val after : ('a,'b) t -> 'a -> ('a * 'b) iter
|
||||
val before : ('a,'b) t -> 'a -> ('a * 'b) iter
|
||||
val iter : ('a,'b) t -> ('a * 'b) iter
|
||||
val add : ('a,'b) t -> ('a * 'b) iter -> ('a,'b) t
|
||||
end
|
||||
|
||||
type ('a,'b) explore =
|
||||
| Yield of 'a * 'b
|
||||
| Explore of ('a, 'b) t
|
||||
| Explore of ('a, 'b) tree
|
||||
|
||||
exception EndOfIter
|
||||
|
||||
|
|
@ -338,16 +364,18 @@ module KList = struct
|
|||
let k, v, l' = _pop l in
|
||||
`Cons ((k,v), _next l')
|
||||
|
||||
let iter t = _next (_push t []) ()
|
||||
let iter {t; _} = _next (_push t []) ()
|
||||
|
||||
let rec add ~cmp t (l:'a t) = match l with
|
||||
let rec _add ~cmp t (l:'a t) = match l with
|
||||
| `Nil -> t
|
||||
| `Cons ((k,v), l') ->
|
||||
add ~cmp (insert ~cmp t k v) (l' ())
|
||||
_add ~cmp (_insert ~cmp t k v) (l' ())
|
||||
|
||||
let after ~cmp t key = _next (_after ~cmp [] t key) ()
|
||||
let add {cmp; t} l = {cmp; t=_add ~cmp t l}
|
||||
|
||||
let before ~cmp t key = _next (_before ~cmp [] t key) ()
|
||||
let after {cmp; t} key = _next (_after ~cmp [] t key) ()
|
||||
|
||||
let before {cmp; t} key = _next (_before ~cmp [] t key) ()
|
||||
end
|
||||
|
||||
module Gen = struct
|
||||
|
|
@ -355,7 +383,7 @@ module Gen = struct
|
|||
|
||||
let _gen stack =
|
||||
let stack = ref stack in
|
||||
let rec next () =
|
||||
let next () =
|
||||
match !stack with
|
||||
| [] -> None
|
||||
| l ->
|
||||
|
|
@ -364,13 +392,15 @@ module Gen = struct
|
|||
Some (k, v)
|
||||
in next
|
||||
|
||||
let iter t = _gen (_push t [])
|
||||
let iter {t; _} = _gen (_push t [])
|
||||
|
||||
let rec add ~cmp t gen =
|
||||
let rec _add ~cmp t gen =
|
||||
match gen() with
|
||||
| None -> t
|
||||
| Some (k,v) -> add ~cmp (insert ~cmp t k v) gen
|
||||
| Some (k,v) -> _add ~cmp (_insert ~cmp t k v) gen
|
||||
|
||||
let after ~cmp t key = _gen (_after ~cmp [] t key)
|
||||
let before ~cmp t key = _gen (_before ~cmp [] t key)
|
||||
let add {cmp; t} l = {cmp; t=_add ~cmp t l}
|
||||
|
||||
let after {cmp; t} key = _gen (_after ~cmp [] t key)
|
||||
let before {cmp; t} key = _gen (_before ~cmp [] t key)
|
||||
end
|
||||
|
|
|
|||
36
misc/AVL.mli
36
misc/AVL.mli
|
|
@ -26,16 +26,20 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
|||
|
||||
(** {1 AVL trees} *)
|
||||
|
||||
type ('a,'b) t = private
|
||||
| Empty
|
||||
| Node of ('a,'b) t * 'a * 'b * ('a,'b) t * int
|
||||
|
||||
type 'a comparator = 'a -> 'a -> int
|
||||
|
||||
val empty : ('a,'b) t
|
||||
type ('a,'b) tree = private
|
||||
| Empty
|
||||
| Node of ('a,'b) tree * 'a * 'b * ('a,'b) tree * int
|
||||
|
||||
type ('a,'b) t = private
|
||||
{ cmp: 'a comparator;
|
||||
t: ('a,'b) tree }
|
||||
|
||||
val empty : cmp:'a comparator -> ('a,'b) t
|
||||
(** Empty tree *)
|
||||
|
||||
val singleton : 'a -> 'b -> ('a,'b) t
|
||||
val singleton : cmp:'a comparator -> 'a -> 'b -> ('a,'b) t
|
||||
(** Tree with a single node *)
|
||||
|
||||
val fold : ('c -> 'a -> 'b -> 'c) -> 'c -> ('a,'b) t -> 'c
|
||||
|
|
@ -44,29 +48,29 @@ val fold : ('c -> 'a -> 'b -> 'c) -> 'c -> ('a,'b) t -> 'c
|
|||
val for_all : ('a -> 'b -> bool) -> ('a,'b) t -> bool
|
||||
val exists : ('a -> 'b -> bool) -> ('a,'b) t -> bool
|
||||
|
||||
val find : cmp:'a comparator -> ('a,'b) t -> 'a -> 'b option
|
||||
val find : ('a,'b) t -> 'a -> 'b option
|
||||
(** Find the value associated to the key, if any *)
|
||||
|
||||
val find_exn : cmp:'a comparator -> ('a,'b) t -> 'a -> 'b
|
||||
val find_exn : ('a,'b) t -> 'a -> 'b
|
||||
(** @raise Not_found if the key is not present *)
|
||||
|
||||
val insert : cmp:'a comparator -> ('a,'b) t -> 'a -> 'b -> ('a,'b) t
|
||||
val insert : ('a,'b) t -> 'a -> 'b -> ('a,'b) t
|
||||
(** Insertion in the tree *)
|
||||
|
||||
val remove : cmp:'a comparator -> ('a,'b) t -> 'a -> ('a,'b) t
|
||||
val remove : ('a,'b) t -> 'a -> ('a,'b) t
|
||||
(** Removal from the tree *)
|
||||
|
||||
val update : cmp:'a comparator -> ('a,'b) t -> 'a ->
|
||||
val update : ('a,'b) t -> 'a ->
|
||||
('b option -> ('a * 'b) option) -> ('a,'b) t
|
||||
(** Update of the given key binding (subsumes [insert] and [remove]) *)
|
||||
|
||||
val split : cmp:'a comparator -> ('a,'b) t -> 'a ->
|
||||
val split : ('a,'b) t -> 'a ->
|
||||
('a,'b) t * 'b option * ('a,'b) t
|
||||
(** [split ~cmp t k] splits [t] into a left part that
|
||||
is smaller than [k], the possible binding of [k],
|
||||
and a part bigger than [k]. *)
|
||||
|
||||
val merge : cmp:'a comparator ->
|
||||
val merge :
|
||||
('a -> 'b option -> 'c option -> 'd option) ->
|
||||
('a,'b) t -> ('a,'c) t -> ('a,'d) t
|
||||
(** Merge two trees together, with the given function *)
|
||||
|
|
@ -82,10 +86,10 @@ val to_list : ('a,'b) t -> ('a * 'b) list
|
|||
module type ITERATOR = sig
|
||||
type 'a iter
|
||||
|
||||
val after : cmp:'a comparator -> ('a,'b) t -> 'a -> ('a * 'b) iter
|
||||
val before : cmp:'a comparator -> ('a,'b) t -> 'a -> ('a * 'b) iter
|
||||
val after : ('a,'b) t -> 'a -> ('a * 'b) iter
|
||||
val before : ('a,'b) t -> 'a -> ('a * 'b) iter
|
||||
val iter : ('a,'b) t -> ('a * 'b) iter
|
||||
val add : cmp:'a comparator -> ('a,'b) t -> ('a * 'b) iter -> ('a,'b) t
|
||||
val add : ('a,'b) t -> ('a * 'b) iter -> ('a,'b) t
|
||||
end
|
||||
|
||||
module KList : sig
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue