src/utils/lib/set2.ml

   1 (***********************************************************************)
   2 (*                                                                     *)
   3 (*                           Objective Caml                            *)
   4 (*                                                                     *)
   5 (*            Xavier Leroy, projet Cristal, INRIA Rocquencourt         *)
   6 (*                                                                     *)
   7 (*  Copyright 1996 Institut National de Recherche en Informatique et   *)
   8 (*  en Automatique.  All rights reserved.  This file is distributed    *)
   9 (*  under the terms of the GNU Library General Public License, with    *)
  10 (*  the special exception on linking described in file ../LICENSE.     *)
  11 (*                                                                     *)
  12 (***********************************************************************)
  13
  14 (* $Id$ *)
  15
  16 (* Sets over ordered types *)
  17
  18 module type OrderedType =
  19   sig
  20     type t
  21     val compare: t -> t -> int
  22   end
  23
  24 module type S =
  25   sig
  26     type elt
  27     type t
  28     val empty: t
  29     val is_empty: t -> bool
  30     val mem: elt -> t -> bool
  31     val add: elt -> t -> t
  32     val singleton: elt -> t
  33     val remove: elt -> t -> t
  34     val union: t -> t -> t
  35     val inter: t -> t -> t
  36     val diff: t -> t -> t
  37     val compare: t -> t -> int
  38     val equal: t -> t -> bool
  39     val subset: t -> t -> bool
  40     val iter: (elt -> unit) -> t -> unit
  41     val fold: (elt -> 'a -> 'a) -> t -> 'a -> 'a
  42     val for_all: (elt -> bool) -> t -> bool
  43     val exists: (elt -> bool) -> t -> bool
  44     val filter: (elt -> bool) -> t -> t
  45     val partition: (elt -> bool) -> t -> t * t
  46     val cardinal: t -> int
  47     val elements: t -> elt list
  48     val min_elt: t -> elt
  49     val max_elt: t -> elt
  50     val choose: t -> elt
  51     val split: elt -> t -> t * bool * t
  52   end
  53
  54 module Make(Ord: OrderedType) =
  55   struct
  56     type elt = Ord.t
  57     type btree = Empty | Node of btree * elt * btree * int
  58     type t = {
  59         tree : btree;
  60         card : int Lazy.t
  61       }
  62
  63     (* Sets are represented by balanced binary trees (the heights of the
  64        children differ by at most 2 *)
  65
  66     let height = function
  67         Empty -> 0
  68       | Node(_, _, _, h) -> h
  69
  70     let rec cardinal_tree = function
  71         Empty -> 0
  72       | Node(l, v, r, _) -> cardinal_tree l + 1 + cardinal_tree r
  73
  74     let t_of_tree t =
  75       { tree = t;
  76         card = lazy (cardinal_tree t); }
  77
  78     let cardinal t =
  79       Lazy.force t.card
  80
  81     (* Creates a new node with left son l, value v and right son r.
  82        We must have all elements of l < v < all elements of r.
  83        l and r must be balanced and | height l - height r | <= 2.
  84        Inline expansion of height for better speed. *)
  85
  86     let create l v r =
  87       let hl = match l with Empty -> 0 | Node(_,_,_,h) -> h in
  88       let hr = match r with Empty -> 0 | Node(_,_,_,h) -> h in
  89       Node(l, v, r, (if hl >= hr then hl + 1 else hr + 1))
  90
  91     (* Same as create, but performs one step of rebalancing if necessary.
  92        Assumes l and r balanced and | height l - height r | <= 3.
  93        Inline expansion of create for better speed in the most frequent case
  94        where no rebalancing is required. *)
  95
  96     let bal l v r =
  97       let hl = match l with Empty -> 0 | Node(_,_,_,h) -> h in
  98       let hr = match r with Empty -> 0 | Node(_,_,_,h) -> h in
  99       if hl > hr + 2 then begin
 100         match l with
 101           Empty -> invalid_arg "Set2.bal"
 102         | Node(ll, lv, lr, _) ->
 103             if height ll >= height lr then
 104               create ll lv (create lr v r)
 105             else begin
 106               match lr with
 107                 Empty -> invalid_arg "Set2.bal"
 108               | Node(lrl, lrv, lrr, _)->
 109                   create (create ll lv lrl) lrv (create lrr v r)
 110             end
 111       end else if hr > hl + 2 then begin
 112         match r with
 113           Empty -> invalid_arg "Set2.bal"
 114         | Node(rl, rv, rr, _) ->
 115             if height rr >= height rl then
 116               create (create l v rl) rv rr
 117             else begin
 118               match rl with
 119                 Empty -> invalid_arg "Set2.bal"
 120               | Node(rll, rlv, rlr, _) ->
 121                   create (create l v rll) rlv (create rlr rv rr)
 122             end
 123       end else
 124         Node(l, v, r, (if hl >= hr then hl + 1 else hr + 1))
 125
 126     (* Insertion of one element *)
 127
 128     let rec add_tree x = function
 129         Empty -> Node(Empty, x, Empty, 1)
 130       | Node(l, v, r, _) as t ->
 131           let c = Ord.compare x v in
 132           if c = 0 then t else
 133           if c < 0 then bal (add_tree x l) v r else bal l v (add_tree x r)
 134
 135     let add x t =
 136       let tree = add_tree x t.tree in
 137       {
 138        tree = tree;
 139        card = if Lazy.lazy_is_val t.card then
 140          Lazy.lazy_from_val ((Lazy.force_val t.card) + 1)
 141        else
 142          lazy (cardinal_tree tree)
 143      }
 144
 145     (* Same as create and bal, but no assumptions are made on the
 146        relative heights of l and r. *)
 147
 148     let rec join l v r =
 149       match (l, r) with
 150         (Empty, _) -> add_tree v r
 151       | (_, Empty) -> add_tree v l
 152       | (Node(ll, lv, lr, lh), Node(rl, rv, rr, rh)) ->
 153           if lh > rh + 2 then bal ll lv (join lr v r) else
 154           if rh > lh + 2 then bal (join l v rl) rv rr else
 155           create l v r
 156
 157     (* Smallest and greatest element of a set *)
 158
 159     let rec min_elt_tree = function
 160         Empty -> raise Not_found
 161       | Node(Empty, v, r, _) -> v
 162       | Node(l, v, r, _) -> min_elt_tree l
 163
 164     let min_elt t =
 165       min_elt_tree t.tree
 166
 167     let rec max_elt_tree = function
 168         Empty -> raise Not_found
 169       | Node(l, v, Empty, _) -> v
 170       | Node(l, v, r, _) -> max_elt_tree r
 171
 172     let max_elt t =
 173       max_elt_tree t.tree
 174
 175     (* Remove the smallest element of the given set *)
 176
 177     let rec remove_min_elt = function
 178         Empty -> invalid_arg "Set2.remove_min_elt"
 179       | Node(Empty, v, r, _) -> r
 180       | Node(l, v, r, _) -> bal (remove_min_elt l) v r
 181
 182     (* Merge two trees l and r into one.
 183        All elements of l must precede the elements of r.
 184        Assume | height l - height r | <= 2. *)
 185
 186     let merge t1 t2 =
 187       match (t1, t2) with
 188         (Empty, t) -> t
 189       | (t, Empty) -> t
 190       | (_, _) -> bal t1 (min_elt_tree t2) (remove_min_elt t2)
 191
 192     (* Merge two trees l and r into one.
 193        All elements of l must precede the elements of r.
 194        No assumption on the heights of l and r. *)
 195
 196     let concat t1 t2 =
 197       match (t1, t2) with
 198         (Empty, t) -> t
 199       | (t, Empty) -> t
 200       | (_, _) -> join t1 (min_elt_tree t2) (remove_min_elt t2)
 201
 202     (* Splitting.  split x s returns a triple (l, present, r) where
 203         - l is the set of elements of s that are < x
 204         - r is the set of elements of s that are > x
 205         - present is false if s contains no element equal to x,
 206           or true if s contains an element equal to x. *)
 207
 208     let rec split_tree x = function
 209         Empty ->
 210           (Empty, false, Empty)
 211       | Node(l, v, r, _) ->
 212           let c = Ord.compare x v in
 213           if c = 0 then (l, true, r)
 214           else if c < 0 then
 215             let (ll, pres, rl) = split_tree x l in (ll, pres, join rl v r)
 216           else
 217             let (lr, pres, rr) = split_tree x r in (join l v lr, pres, rr)
 218
 219     let split e t =
 220       let (ll, pres, rl) = split_tree e t.tree in
 221       (t_of_tree ll, pres, t_of_tree rl)
 222
 223     (* Implementation of the set operations *)
 224
 225     let empty = { tree = Empty; card = Lazy.lazy_from_val 0; }
 226
 227     let is_empty_tree = function Empty -> true | _ -> false
 228
 229     let is_empty t = is_empty_tree t.tree
 230
 231     let rec mem_tree x = function
 232         Empty -> false
 233       | Node(l, v, r, _) ->
 234           let c = Ord.compare x v in
 235           c = 0 || mem_tree x (if c < 0 then l else r)
 236
 237     let mem x t = mem_tree x t.tree
 238
 239     let singleton x =
 240       { tree = Node(Empty, x, Empty, 1);
 241         card = Lazy.lazy_from_val 1; }
 242
 243     let rec remove_tree x = function
 244         Empty -> (Empty, false)
 245       | Node(l, v, r, _) ->
 246           let c = Ord.compare x v in
 247           if c = 0 then (merge l r, true) else
 248           if c < 0 then
 249             let tree, found = remove_tree x l in
 250             (bal tree v r, found)
 251           else
 252             let tree, found = remove_tree x r in
 253             (bal l v tree, found)
 254
 255     let remove x t =
 256       let tree, found = remove_tree x t.tree in
 257       if found then
 258         { tree = tree;
 259           card = if Lazy.lazy_is_val t.card then
 260             Lazy.lazy_from_val ((Lazy.force_val t.card) - 1)
 261           else
 262             lazy (cardinal_tree tree) }
 263       else t
 264
 265     let rec union_tree s1 s2 =
 266       match (s1, s2) with
 267         (Empty, t2) -> t2
 268       | (t1, Empty) -> t1
 269       | (Node(l1, v1, r1, h1), Node(l2, v2, r2, h2)) ->
 270           if h1 >= h2 then
 271             if h2 = 1 then add_tree v2 s1 else begin
 272               let (l2, _, r2) = split_tree v1 s2 in
 273               join (union_tree l1 l2) v1 (union_tree r1 r2)
 274             end
 275           else
 276             if h1 = 1 then add_tree v1 s2 else begin
 277               let (l1, _, r1) = split_tree v2 s1 in
 278               join (union_tree l1 l2) v2 (union_tree r1 r2)
 279             end
 280
 281     let union s1 s2 =
 282       t_of_tree (union_tree s1.tree s2.tree)
 283
 284     let rec inter_tree s1 s2 =
 285       match (s1, s2) with
 286         (Empty, t2) -> Empty
 287       | (t1, Empty) -> Empty
 288       | (Node(l1, v1, r1, _), t2) ->
 289           match split_tree v1 t2 with
 290             (l2, false, r2) ->
 291               concat (inter_tree l1 l2) (inter_tree r1 r2)
 292           | (l2, true, r2) ->
 293               join (inter_tree l1 l2) v1 (inter_tree r1 r2)
 294
 295     let inter s1 s2 =
 296       t_of_tree (inter_tree s1.tree s2.tree)
 297
 298     let rec diff_tree s1 s2 =
 299       match (s1, s2) with
 300         (Empty, t2) -> Empty
 301       | (t1, Empty) -> t1
 302       | (Node(l1, v1, r1, _), t2) ->
 303           match split_tree v1 t2 with
 304             (l2, false, r2) ->
 305               join (diff_tree l1 l2) v1 (diff_tree r1 r2)
 306           | (l2, true, r2) ->
 307               concat (diff_tree l1 l2) (diff_tree r1 r2)
 308
 309     let diff s1 s2 =
 310       t_of_tree (diff_tree s1.tree s2.tree)
 311
 312     type enumeration = End | More of elt * btree * enumeration
 313
 314     let rec cons_enum s e =
 315       match s with
 316         Empty -> e
 317       | Node(l, v, r, _) -> cons_enum l (More(v, r, e))
 318
 319     let rec compare_aux e1 e2 =
 320         match (e1, e2) with
 321         (End, End) -> 0
 322       | (End, _)  -> -1
 323       | (_, End) -> 1
 324       | (More(v1, r1, e1), More(v2, r2, e2)) ->
 325           let c = Ord.compare v1 v2 in
 326           if c <> 0
 327           then c
 328           else compare_aux (cons_enum r1 e1) (cons_enum r2 e2)
 329
 330     let compare_tree s1 s2 =
 331       compare_aux (cons_enum s1 End) (cons_enum s2 End)
 332
 333     let compare s1 s2 =
 334       compare_tree s1.tree s2.tree
 335
 336     let equal s1 s2 =
 337       compare s1 s2 = 0
 338
 339     let rec subset_tree s1 s2 =
 340       match (s1, s2) with
 341         Empty, _ ->
 342           true
 343       | _, Empty ->
 344           false
 345       | Node (l1, v1, r1, _), (Node (l2, v2, r2, _) as t2) ->
 346           let c = Ord.compare v1 v2 in
 347           if c = 0 then
 348             subset_tree l1 l2 && subset_tree r1 r2
 349           else if c < 0 then
 350             subset_tree (Node (l1, v1, Empty, 0)) l2 && subset_tree r1 t2
 351           else
 352             subset_tree (Node (Empty, v1, r1, 0)) r2 && subset_tree l1 t2
 353
 354     let subset s1 s2 =
 355       subset_tree s1.tree s2.tree
 356
 357     let rec iter_tree f = function
 358         Empty -> ()
 359       | Node(l, v, r, _) -> iter_tree f l; f v; iter_tree f r
 360
 361     let iter f t =
 362       iter_tree f t.tree
 363
 364     let rec fold_tree f s accu =
 365       match s with
 366         Empty -> accu
 367       | Node(l, v, r, _) -> fold_tree f r (f v (fold_tree f l accu))
 368
 369     let fold f s accu =
 370       fold_tree f s.tree accu
 371
 372     let rec for_all_tree p = function
 373         Empty -> true
 374       | Node(l, v, r, _) -> p v && for_all_tree p l && for_all_tree p r
 375
 376     let for_all p t =
 377       for_all_tree p t.tree
 378
 379     let rec exists_tree p = function
 380         Empty -> false
 381       | Node(l, v, r, _) -> p v || exists_tree p l || exists_tree p r
 382
 383     let exists p t =
 384       exists_tree p t.tree
 385
 386     let filter_tree p s =
 387       let rec filt accu = function
 388         | Empty -> accu
 389         | Node(l, v, r, _) ->
 390             filt (filt (if p v then add_tree v accu else accu) l) r in
 391       filt Empty s
 392
 393     let filter p s =
 394       t_of_tree (filter_tree p s.tree)
 395
 396     let partition_tree p s =
 397       let rec part (t, f as accu) = function
 398         | Empty -> accu
 399         | Node(l, v, r, _) ->
 400             part (part (if p v then (add_tree v t, f) else (t, add_tree v f)) l) r in
 401       part (Empty, Empty) s
 402
 403     let partition p s =
 404       let tree1, tree2 = partition_tree p s.tree in
 405       (t_of_tree tree1, t_of_tree tree2)
 406
 407     let rec elements_aux accu = function
 408         Empty -> accu
 409       | Node(l, v, r, _) -> elements_aux (v :: elements_aux accu r) l
 410
 411     let elements s =
 412       elements_aux [] s.tree
 413
 414     let choose = min_elt
 415
 416   end