227 lines
6.7 KiB
OCaml
227 lines
6.7 KiB
OCaml
(* Sets over ordered types *)
|
|
|
|
module type OrderedType =
|
|
sig
|
|
type t
|
|
val compare: t -> t -> int
|
|
end
|
|
|
|
module type S =
|
|
sig
|
|
type elt
|
|
type t
|
|
val empty: t
|
|
val is_empty: t -> bool
|
|
val mem: elt -> t -> bool
|
|
val add: elt -> t -> t
|
|
val remove: elt -> t -> t
|
|
val union: t -> t -> t
|
|
val inter: t -> t -> t
|
|
val diff: t -> t -> t
|
|
val compare: t -> t -> int
|
|
val equal: t -> t -> bool
|
|
val iter: (elt -> 'a) -> t -> unit
|
|
val fold: (elt -> 'a -> 'a) -> t -> 'a -> 'a
|
|
val elements: t -> elt list
|
|
val choose: t -> elt
|
|
end
|
|
|
|
module Make(Ord: OrderedType): (S with elt = Ord.t) =
|
|
struct
|
|
type elt = Ord.t
|
|
type t = Empty | Node of t * elt * t * int
|
|
|
|
(* Sets are represented by balanced binary trees (the heights of the
|
|
children differ by at most 2 *)
|
|
|
|
let height = function
|
|
Empty -> 0
|
|
| Node(_, _, _, h) -> h
|
|
|
|
(* Creates a new node with left son l, value x and right son r.
|
|
l and r must be balanced and | height l - height r | <= 2.
|
|
Inline expansion of height for better speed. *)
|
|
|
|
let new l x r =
|
|
let hl = match l with Empty -> 0 | Node(_,_,_,h) -> h in
|
|
let hr = match r with Empty -> 0 | Node(_,_,_,h) -> h in
|
|
Node(l, x, r, (if hl >= hr then hl + 1 else hr + 1))
|
|
|
|
(* Same as new, but performs one step of rebalancing if necessary.
|
|
Assumes l and r balanced.
|
|
Inline expansion of new for better speed in the most frequent case
|
|
where no rebalancing is required. *)
|
|
|
|
let bal l x r =
|
|
let hl = match l with Empty -> 0 | Node(_,_,_,h) -> h in
|
|
let hr = match r with Empty -> 0 | Node(_,_,_,h) -> h in
|
|
if hl > hr + 2 then begin
|
|
match l with
|
|
Empty -> invalid_arg "Set.bal"
|
|
| Node(ll, lv, lr, _) ->
|
|
if height ll >= height lr then
|
|
new ll lv (new lr x r)
|
|
else begin
|
|
match lr with
|
|
Empty -> invalid_arg "Set.bal"
|
|
| Node(lrl, lrv, lrr, _)->
|
|
new (new ll lv lrl) lrv (new lrr x r)
|
|
end
|
|
end else if hr > hl + 2 then begin
|
|
match r with
|
|
Empty -> invalid_arg "Set.bal"
|
|
| Node(rl, rv, rr, _) ->
|
|
if height rr >= height rl then
|
|
new (new l x rl) rv rr
|
|
else begin
|
|
match rl with
|
|
Empty -> invalid_arg "Set.bal"
|
|
| Node(rll, rlv, rlr, _) ->
|
|
new (new l x rll) rlv (new rlr rv rr)
|
|
end
|
|
end else
|
|
Node(l, x, r, (if hl >= hr then hl + 1 else hr + 1))
|
|
|
|
(* Same as bal, but repeat rebalancing until the final result
|
|
is balanced. *)
|
|
|
|
let rec join l x r =
|
|
match bal l x r with
|
|
Empty -> invalid_arg "Set.join"
|
|
| Node(l', x', r', _) as t' ->
|
|
let d = height l' - height r' in
|
|
if d < -2 or d > 2 then join l' x' r' else t'
|
|
|
|
(* Merge two trees l and r into one.
|
|
All elements of l must precede the elements of r.
|
|
Assumes | height l - height r | <= 2. *)
|
|
|
|
let rec merge t1 t2 =
|
|
match (t1, t2) with
|
|
(Empty, t) -> t
|
|
| (t, Empty) -> t
|
|
| (Node(l1, v1, r1, h1), Node(l2, v2, r2, h2)) ->
|
|
bal l1 v1 (bal (merge r1 l2) v2 r2)
|
|
|
|
(* Same as merge, but does not assume anything about l and r. *)
|
|
|
|
let rec concat t1 t2 =
|
|
match (t1, t2) with
|
|
(Empty, t) -> t
|
|
| (t, Empty) -> t
|
|
| (Node(l1, v1, r1, h1), Node(l2, v2, r2, h2)) ->
|
|
join l1 v1 (join (concat r1 l2) v2 r2)
|
|
|
|
(* Splitting *)
|
|
|
|
let rec split x = function
|
|
Empty ->
|
|
(Empty, None, Empty)
|
|
| Node(l, v, r, _) ->
|
|
let c = Ord.compare x v in
|
|
if c = 0 then (l, Some v, r)
|
|
else if c < 0 then
|
|
let (ll, vl, rl) = split x l in (ll, vl, join rl v r)
|
|
else
|
|
let (lr, vr, rr) = split x r in (join l v lr, vr, rr)
|
|
|
|
(* Implementation of the set operations *)
|
|
|
|
let empty = Empty
|
|
|
|
let is_empty = function Empty -> true | _ -> false
|
|
|
|
let rec mem x = function
|
|
Empty -> false
|
|
| Node(l, v, r, _) ->
|
|
let c = Ord.compare x v in
|
|
if c = 0 then true else
|
|
if c < 0 then mem x l else mem x r
|
|
|
|
let rec add x = function
|
|
Empty -> Node(Empty, x, Empty, 1)
|
|
| Node(l, v, r, _) as t ->
|
|
let c = Ord.compare x v in
|
|
if c = 0 then t else
|
|
if c < 0 then bal (add x l) v r else bal l v (add x r)
|
|
|
|
let rec remove x = function
|
|
Empty -> Empty
|
|
| Node(l, v, r, _) ->
|
|
let c = Ord.compare x v in
|
|
if c = 0 then merge l r else
|
|
if c < 0 then bal (remove x l) v r else bal l v (remove x r)
|
|
|
|
let rec union s1 s2 =
|
|
match (s1, s2) with
|
|
(Empty, t2) -> t2
|
|
| (t1, Empty) -> t1
|
|
| (Node(l1, v1, r1, _), t2) ->
|
|
let (l2, _, r2) = split v1 t2 in
|
|
join (union l1 l2) v1 (union r1 r2)
|
|
|
|
let rec inter s1 s2 =
|
|
match (s1, s2) with
|
|
(Empty, t2) -> Empty
|
|
| (t1, Empty) -> Empty
|
|
| (Node(l1, v1, r1, _), t2) ->
|
|
match split v1 t2 with
|
|
(l2, None, r2) ->
|
|
concat (inter l1 l2) (inter r1 r2)
|
|
| (l2, Some _, r2) ->
|
|
join (inter l1 l2) v1 (inter r1 r2)
|
|
|
|
let rec diff s1 s2 =
|
|
match (s1, s2) with
|
|
(Empty, t2) -> Empty
|
|
| (t1, Empty) -> t1
|
|
| (Node(l1, v1, r1, _), t2) ->
|
|
match split v1 t2 with
|
|
(l2, None, r2) ->
|
|
join (diff l1 l2) v1 (diff r1 r2)
|
|
| (l2, Some _, r2) ->
|
|
concat (diff l1 l2) (diff r1 r2)
|
|
|
|
let rec compare_aux l1 l2 =
|
|
match (l1, l2) with
|
|
([], []) -> 0
|
|
| ([], _) -> -1
|
|
| (_, []) -> 1
|
|
| (Empty :: t1, Empty :: t2) ->
|
|
compare_aux t1 t2
|
|
| (Node(Empty, v1, r1, _) :: t1, Node(Empty, v2, r2, _) :: t2) ->
|
|
let c = Ord.compare v1 v2 in
|
|
if c <> 0 then c else compare_aux (r1::t1) (r2::t2)
|
|
| (Node(l1, v1, r1, _) :: t1, t2) ->
|
|
compare_aux (l1 :: Node(Empty, v1, r1, 0) :: t1) t2
|
|
| (t1, Node(l2, v2, r2, _) :: t2) ->
|
|
compare_aux t1 (l2 :: Node(Empty, v2, r2, 0) :: t2)
|
|
|
|
let compare s1 s2 =
|
|
compare_aux [s1] [s2]
|
|
|
|
let equal s1 s2 =
|
|
compare s1 s2 = 0
|
|
|
|
let rec iter f = function
|
|
Empty -> ()
|
|
| Node(l, v, r, _) -> iter f l; f v; iter f r
|
|
|
|
let rec fold f s accu =
|
|
match s with
|
|
Empty -> accu
|
|
| Node(l, v, r, _) -> fold f l (f v (fold f r accu))
|
|
|
|
let rec elements_aux accu = function
|
|
Empty -> accu
|
|
| Node(l, v, r, _) -> elements_aux (v :: elements_aux accu r) l
|
|
|
|
let elements s =
|
|
elements_aux [] s
|
|
|
|
let rec choose = function
|
|
Empty -> raise Not_found
|
|
| Node(Empty, v, r, _) -> v
|
|
| Node(l, v, r, _) -> choose l
|
|
end
|