ocaml-containers/core/CCMultiSet.ml

(*
copyright (c) 2013, simon cruanes
all rights reserved.

redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

redistributions of source code must retain the above copyright notice, this
list of conditions and the following disclaimer.  redistributions in binary
form must reproduce the above copyright notice, this list of conditions and the
following disclaimer in the documentation and/or other materials provided with
the distribution.

this software is provided by the copyright holders and contributors "as is" and
any express or implied warranties, including, but not limited to, the implied
warranties of merchantability and fitness for a particular purpose are
disclaimed. in no event shall the copyright holder or contributors be liable
for any direct, indirect, incidental, special, exemplary, or consequential
  damages (including, but not limited to, procurement of substitute goods or
  services; loss of use, data, or profits; or business interruption) however
  caused and on any theory of liability, whether in contract, strict liability,
  or tort (including negligence or otherwise) arising in any way out of the use
  of this software, even if advised of the possibility of such damage.
*)

(** {1 Multiset} *)

type 'a sequence = ('a -> unit) -> unit

module type S = sig
  type elt
  type t

  val empty : t

  val is_empty : t -> bool

  val mem : t -> elt -> bool

  val count : t -> elt -> int

  val singleton : elt -> t

  val add : t -> elt -> t

  val remove : t -> elt -> t

  val add_mult : t -> elt -> int -> t

  val remove_mult : t -> elt -> int -> t

  val update : t -> elt -> (int -> int) -> t

  val min : t -> elt

  val max : t -> elt

  val union : t -> t -> t

  val meet : t -> t -> t

  val intersection : t -> t -> t

  val diff : t -> t -> t

  val contains : t -> t -> bool

  val compare : t -> t -> int

  val equal : t -> t -> bool

  val cardinal : t -> int
    (** Number of distinct elements *)

  val iter : t -> (int -> elt -> unit) -> unit

  val fold : t -> 'b -> ('b -> int -> elt -> 'b) -> 'b

  val of_list : elt list -> t

  val to_list : t -> elt list

  val to_seq : t -> elt sequence

  val of_seq : elt sequence -> t
end

module Make(O : Set.OrderedType) = struct
  module M = Map.Make(O)

  type t = int M.t

  type elt = O.t

  let empty = M.empty

  let is_empty = M.is_empty

  let mem ms x = M.mem x ms

  let count ms x =
    try M.find x ms
    with Not_found -> 0

  let singleton x = M.singleton x 1

  let add ms x =
    let n = count ms x in
    M.add x (n+1) ms

  let add_mult ms x n =
    if n < 0 then invalid_arg "CCMultiSet.add_mult";
    if n=0
      then ms
      else M.add x (count ms x + n) ms

  let remove_mult ms x n =
    if n < 0 then invalid_arg "CCMultiSet.remove_mult";
    let cur_n = count ms x in
    let new_n = cur_n - n in
    if new_n <= 0
      then M.remove x ms
      else M.add x new_n ms

  let remove ms x = remove_mult ms x 1

  let update ms x f =
    let n = count ms x in
    match f n with
    | 0 ->
        if n=0 then ms else M.remove x ms
    | n' ->
        if n' < 0
          then invalid_arg "CCMultiSet.udpate"
          else M.add x n' ms

  let min ms =
    fst (M.min_binding ms)

  let max ms =
    fst (M.max_binding ms)

  let union m1 m2 =
    M.merge
      (fun _x n1 n2 -> match n1, n2 with
        | None, None -> assert false
        | Some n, None
        | None, Some n -> Some n
        | Some n1, Some n2 -> Some (n1+n2))
      m1 m2

  let meet m1 m2 =
	M.merge
	  (fun _ n1 n2 -> match n1, n2 with
		| None, None -> assert false
		| Some n, None | None, Some n -> Some n
		| Some n1, Some n2 -> Some (Pervasives.max n1 n2))
	  m1 m2

  let intersection m1 m2 =
    M.merge
      (fun _x n1 n2 -> match n1, n2 with
        | None, None -> assert false
        | Some _, None
        | None, Some _ -> None
        | Some n1, Some n2 -> Some (Pervasives.min n1 n2))
      m1 m2

  let diff m1 m2 =
    M.merge
      (fun _x n1 n2 -> match n1, n2 with
        | None, None -> assert false
        | Some n1, None -> Some n1
        | None, Some _n2 -> None
        | Some n1, Some n2 ->
          if n1 > n2
            then Some (n1 - n2)
            else None)
      m1 m2

  let contains m1 m2 =
    try
      M.for_all (fun x c -> M.find x m1 >= c) m2
    with Not_found -> false

  let compare m1 m2 =
    M.compare (fun x y -> x - y) m1 m2

  let equal m1 m2 =
    M.equal (fun x y -> x = y) m1 m2

  let cardinal m = M.cardinal m

  let iter m f =
    M.iter (fun x n -> f n x) m

  let fold m acc f =
    M.fold (fun x n acc -> f acc n x) m acc

  let of_list l =
    let rec build acc l = match l with
    | [] -> acc
    | x::l' -> build (add acc x) l'
    in
    build empty l

  let to_list m =
    (* [n_cons n x l] is the result of applying [fun l -> x :: l]  [n] times
        to [l] *)
    let rec n_cons n x l = match n with
    | 0 -> l
    | 1 -> x::l
    | _ -> n_cons (n-1) x (x::l)
    in
    fold m [] (fun acc n x -> n_cons n x acc)

  let to_seq m k =
    M.iter (fun x n -> for _i = 1 to n do k x done) m

  let of_seq seq =
    let m = ref empty in
    seq (fun x -> m := add !m x);
    !m
end

(*$T
  let module S = CCMultiSet.Make(String) in \
    S.count (S.add_mult S.empty "a" 5) "a" = 5
  let module S = CCMultiSet.Make(String) in \
    S.count (S.remove_mult (S.add_mult S.empty "a" 5) "a" 3) "a" = 2
  let module S = CCMultiSet.Make(String) in \
    S.count (S.remove_mult (S.add_mult S.empty "a" 4) "a" 6) "a" = 0
*)