l4/pkg/ocaml/contrib/stdlib/set.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: set.ml 6694 2004-11-25 00:06:06Z doligez $ *)
  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 t = Empty | Node of t * elt * t * int
  58
  59     (* Sets are represented by balanced binary trees (the heights of the
  60        children differ by at most 2 *)
  61
  62     let height = function
  63         Empty -> 0
  64       | Node(_, _, _, h) -> h
  65
  66     (* Creates a new node with left son l, value v and right son r.
  67        We must have all elements of l < v < all elements of r.
  68        l and r must be balanced and | height l - height r | <= 2.
  69        Inline expansion of height for better speed. *)
  70
  71     let create l v r =
  72       let hl = match l with Empty -> 0 | Node(_,_,_,h) -> h in
  73       let hr = match r with Empty -> 0 | Node(_,_,_,h) -> h in
  74       Node(l, v, r, (if hl >= hr then hl + 1 else hr + 1))
  75
  76     (* Same as create, but performs one step of rebalancing if necessary.
  77        Assumes l and r balanced and | height l - height r | <= 3.
  78        Inline expansion of create for better speed in the most frequent case
  79        where no rebalancing is required. *)
  80
  81     let bal l v r =
  82       let hl = match l with Empty -> 0 | Node(_,_,_,h) -> h in
  83       let hr = match r with Empty -> 0 | Node(_,_,_,h) -> h in
  84       if hl > hr + 2 then begin
  85         match l with
  86           Empty -> invalid_arg "Set.bal"
  87         | Node(ll, lv, lr, _) ->
  88             if height ll >= height lr then
  89               create ll lv (create lr v r)
  90             else begin
  91               match lr with
  92                 Empty -> invalid_arg "Set.bal"
  93               | Node(lrl, lrv, lrr, _)->
  94                   create (create ll lv lrl) lrv (create lrr v r)
  95             end
  96       end else if hr > hl + 2 then begin
  97         match r with
  98           Empty -> invalid_arg "Set.bal"
  99         | Node(rl, rv, rr, _) ->
 100             if height rr >= height rl then
 101               create (create l v rl) rv rr
 102             else begin
 103               match rl with
 104                 Empty -> invalid_arg "Set.bal"
 105               | Node(rll, rlv, rlr, _) ->
 106                   create (create l v rll) rlv (create rlr rv rr)
 107             end
 108       end else
 109         Node(l, v, r, (if hl >= hr then hl + 1 else hr + 1))
 110
 111     (* Insertion of one element *)
 112
 113     let rec add x = function
 114         Empty -> Node(Empty, x, Empty, 1)
 115       | Node(l, v, r, _) as t ->
 116           let c = Ord.compare x v in
 117           if c = 0 then t else
 118           if c < 0 then bal (add x l) v r else bal l v (add x r)
 119
 120     (* Same as create and bal, but no assumptions are made on the
 121        relative heights of l and r. *)
 122
 123     let rec join l v r =
 124       match (l, r) with
 125         (Empty, _) -> add v r
 126       | (_, Empty) -> add v l
 127       | (Node(ll, lv, lr, lh), Node(rl, rv, rr, rh)) ->
 128           if lh > rh + 2 then bal ll lv (join lr v r) else
 129           if rh > lh + 2 then bal (join l v rl) rv rr else
 130           create l v r
 131
 132     (* Smallest and greatest element of a set *)
 133
 134     let rec min_elt = function
 135         Empty -> raise Not_found
 136       | Node(Empty, v, r, _) -> v
 137       | Node(l, v, r, _) -> min_elt l
 138
 139     let rec max_elt = function
 140         Empty -> raise Not_found
 141       | Node(l, v, Empty, _) -> v
 142       | Node(l, v, r, _) -> max_elt r
 143
 144     (* Remove the smallest element of the given set *)
 145
 146     let rec remove_min_elt = function
 147         Empty -> invalid_arg "Set.remove_min_elt"
 148       | Node(Empty, v, r, _) -> r
 149       | Node(l, v, r, _) -> bal (remove_min_elt l) v r
 150
 151     (* Merge two trees l and r into one.
 152        All elements of l must precede the elements of r.
 153        Assume | height l - height r | <= 2. *)
 154
 155     let merge t1 t2 =
 156       match (t1, t2) with
 157         (Empty, t) -> t
 158       | (t, Empty) -> t
 159       | (_, _) -> bal t1 (min_elt t2) (remove_min_elt t2)
 160
 161     (* Merge two trees l and r into one.
 162        All elements of l must precede the elements of r.
 163        No assumption on the heights of l and r. *)
 164
 165     let concat t1 t2 =
 166       match (t1, t2) with
 167         (Empty, t) -> t
 168       | (t, Empty) -> t
 169       | (_, _) -> join t1 (min_elt t2) (remove_min_elt t2)
 170
 171     (* Splitting.  split x s returns a triple (l, present, r) where
 172         - l is the set of elements of s that are < x
 173         - r is the set of elements of s that are > x
 174         - present is false if s contains no element equal to x,
 175           or true if s contains an element equal to x. *)
 176
 177     let rec split x = function
 178         Empty ->
 179           (Empty, false, Empty)
 180       | Node(l, v, r, _) ->
 181           let c = Ord.compare x v in
 182           if c = 0 then (l, true, r)
 183           else if c < 0 then
 184             let (ll, pres, rl) = split x l in (ll, pres, join rl v r)
 185           else
 186             let (lr, pres, rr) = split x r in (join l v lr, pres, rr)
 187
 188     (* Implementation of the set operations *)
 189
 190     let empty = Empty
 191
 192     let is_empty = function Empty -> true | _ -> false
 193
 194     let rec mem x = function
 195         Empty -> false
 196       | Node(l, v, r, _) ->
 197           let c = Ord.compare x v in
 198           c = 0 || mem x (if c < 0 then l else r)
 199
 200     let singleton x = Node(Empty, x, Empty, 1)
 201
 202     let rec remove x = function
 203         Empty -> Empty
 204       | Node(l, v, r, _) ->
 205           let c = Ord.compare x v in
 206           if c = 0 then merge l r else
 207           if c < 0 then bal (remove x l) v r else bal l v (remove x r)
 208
 209     let rec union s1 s2 =
 210       match (s1, s2) with
 211         (Empty, t2) -> t2
 212       | (t1, Empty) -> t1
 213       | (Node(l1, v1, r1, h1), Node(l2, v2, r2, h2)) ->
 214           if h1 >= h2 then
 215             if h2 = 1 then add v2 s1 else begin
 216               let (l2, _, r2) = split v1 s2 in
 217               join (union l1 l2) v1 (union r1 r2)
 218             end
 219           else
 220             if h1 = 1 then add v1 s2 else begin
 221               let (l1, _, r1) = split v2 s1 in
 222               join (union l1 l2) v2 (union r1 r2)
 223             end
 224
 225     let rec inter s1 s2 =
 226       match (s1, s2) with
 227         (Empty, t2) -> Empty
 228       | (t1, Empty) -> Empty
 229       | (Node(l1, v1, r1, _), t2) ->
 230           match split v1 t2 with
 231             (l2, false, r2) ->
 232               concat (inter l1 l2) (inter r1 r2)
 233           | (l2, true, r2) ->
 234               join (inter l1 l2) v1 (inter r1 r2)
 235
 236     let rec diff s1 s2 =
 237       match (s1, s2) with
 238         (Empty, t2) -> Empty
 239       | (t1, Empty) -> t1
 240       | (Node(l1, v1, r1, _), t2) ->
 241           match split v1 t2 with
 242             (l2, false, r2) ->
 243               join (diff l1 l2) v1 (diff r1 r2)
 244           | (l2, true, r2) ->
 245               concat (diff l1 l2) (diff r1 r2)
 246
 247     type enumeration = End | More of elt * t * enumeration
 248
 249     let rec cons_enum s e =
 250       match s with
 251         Empty -> e
 252       | Node(l, v, r, _) -> cons_enum l (More(v, r, e))
 253
 254     let rec compare_aux e1 e2 =
 255         match (e1, e2) with
 256         (End, End) -> 0
 257       | (End, _)  -> -1
 258       | (_, End) -> 1
 259       | (More(v1, r1, e1), More(v2, r2, e2)) ->
 260           let c = Ord.compare v1 v2 in
 261           if c <> 0
 262           then c
 263           else compare_aux (cons_enum r1 e1) (cons_enum r2 e2)
 264
 265     let compare s1 s2 =
 266       compare_aux (cons_enum s1 End) (cons_enum s2 End)
 267
 268     let equal s1 s2 =
 269       compare s1 s2 = 0
 270
 271     let rec subset s1 s2 =
 272       match (s1, s2) with
 273         Empty, _ ->
 274           true
 275       | _, Empty ->
 276           false
 277       | Node (l1, v1, r1, _), (Node (l2, v2, r2, _) as t2) ->
 278           let c = Ord.compare v1 v2 in
 279           if c = 0 then
 280             subset l1 l2 && subset r1 r2
 281           else if c < 0 then
 282             subset (Node (l1, v1, Empty, 0)) l2 && subset r1 t2
 283           else
 284             subset (Node (Empty, v1, r1, 0)) r2 && subset l1 t2
 285
 286     let rec iter f = function
 287         Empty -> ()
 288       | Node(l, v, r, _) -> iter f l; f v; iter f r
 289
 290     let rec fold f s accu =
 291       match s with
 292         Empty -> accu
 293       | Node(l, v, r, _) -> fold f r (f v (fold f l accu))
 294
 295     let rec for_all p = function
 296         Empty -> true
 297       | Node(l, v, r, _) -> p v && for_all p l && for_all p r
 298
 299     let rec exists p = function
 300         Empty -> false
 301       | Node(l, v, r, _) -> p v || exists p l || exists p r
 302
 303     let filter p s =
 304       let rec filt accu = function
 305         | Empty -> accu
 306         | Node(l, v, r, _) ->
 307             filt (filt (if p v then add v accu else accu) l) r in
 308       filt Empty s
 309
 310     let partition p s =
 311       let rec part (t, f as accu) = function
 312         | Empty -> accu
 313         | Node(l, v, r, _) ->
 314             part (part (if p v then (add v t, f) else (t, add v f)) l) r in
 315       part (Empty, Empty) s
 316
 317     let rec cardinal = function
 318         Empty -> 0
 319       | Node(l, v, r, _) -> cardinal l + 1 + cardinal r
 320
 321     let rec elements_aux accu = function
 322         Empty -> accu
 323       | Node(l, v, r, _) -> elements_aux (v :: elements_aux accu r) l
 324
 325     let elements s =
 326       elements_aux [] s
 327
 328     let choose = min_elt
 329
 330   end