Merge pull request #8 from dlesbre/hset-functions

backtracking · web-flow · commit 4285c88e8928 · 2024-01-28T20:37:38.000+01:00
Extra Hset functions
diff --git a/hashcons.ml b/hashcons.ml
@@ -663,6 +663,9 @@ module Hset = struct
     | Leaf j -> k.tag == j.tag
     | Branch (_, m, l, r) -> mem k (if zero_bit k.tag m then l else r)
 
+  let find k s = if mem k s then k else raise Not_found
+  let find_opt k s = if mem k s then Some k else None
+
   (*s The following operation [join] will be used in both insertion and
       union. Given two non-empty trees [t0] and [t1] with longest common
       prefixes [p0] and [p1] respectively, which are supposed to
@@ -832,8 +835,8 @@ module Hset = struct
 	s1
 
   (*s All the following operations ([cardinal], [iter], [fold], [for_all],
-      [exists], [filter], [partition], [choose], [elements]) are
-      implemented as for any other kind of binary trees. *)
+      [exists], [filter], [partition], [choose], [choose_opt], [elements],
+      [to_seq]) are implemented as for any other kind of binary trees. *)
 
   let rec cardinal = function
     | Empty -> 0
@@ -878,6 +881,11 @@ module Hset = struct
     | Leaf k -> k
     | Branch (_, _,t0,_) -> choose t0   (* we know that [t0] is non-empty *)
 
+  let rec choose_opt = function
+    | Empty -> None
+    | Leaf k -> Some k
+    | Branch (_, _,t0,_) -> choose_opt t0   (* we know that [t0] is non-empty *)
+
   let elements s =
     let rec elements_aux acc = function
       | Empty -> acc
@@ -886,6 +894,36 @@ module Hset = struct
     in
     elements_aux [] s
 
+  let to_seq s =
+    let rec to_seq_aux acc = function
+      | Empty -> acc
+      | Leaf k -> Seq.cons k acc
+      | Branch (_,_,l,r) -> to_seq_aux (to_seq_aux acc r) l
+    in
+    to_seq_aux Seq.empty s
+
+  let split elt s =
+    fold (fun elt' (lt, present, gt) ->
+      if elt'.tag < elt.tag then (add elt' lt, present, gt) else
+      if elt'.tag > elt.tag then (lt, present, add elt' gt) else
+      (lt, true, gt)
+    ) s (Empty, false, Empty)
+
+  (*s [map] and [filter_map] are implemented via [fold] and [add]
+      since we can't relate the tag of [f elt] to that of [elt] *)
+  let map f s = fold (fun elt s -> add (f elt) s) s Empty
+  let filter_map f s = fold (fun elt s ->
+      match f elt with
+        | None -> s
+        | Some elt' -> add elt' s)
+    s Empty
+
+  let add_seq seq s = Seq.fold_left (fun s elt -> add elt s) s seq
+
+  let of_seq seq = add_seq seq Empty
+
+  let of_list list = List.fold_left (fun s elt -> add elt s) Empty list
+
   (*s There is no way to give an efficient implementation of [min_elt]
       and [max_elt], as with binary search trees.  The following
       implementation is a traversal of all elements, barely more
@@ -899,11 +937,53 @@ module Hset = struct
     | Leaf k -> k
     | Branch (_,_,s,t) -> min (min_elt s) (min_elt t)
 
+  let min_elt_opt = function
+    | Empty -> None
+    | x -> Some (min_elt x)
+
   let rec max_elt = function
     | Empty -> raise Not_found
     | Leaf k -> k
     | Branch (_,_,s,t) -> max (max_elt s) (max_elt t)
 
+  let max_elt_opt = function
+    | Empty -> None
+    | x -> Some (max_elt x)
+
+  (*s [find_first], [find_last] and their opt versions are less efficient
+      then with binary search trees. They are linear time and can call [f] an
+      arbitrary number of times, and not necessarily on elements smaller/larger
+      than the witness. *)
+  let find_first_opt f s =
+    fold
+      (fun elt acc ->
+        match acc with
+        | None -> if f elt then Some elt else None
+        | Some witness ->
+            if witness.tag <= elt.tag then acc else
+            if f elt then Some elt else acc)
+      s None
+
+  let find_first f s =
+    match find_first_opt f s with
+      | Some elt -> elt
+      | None -> raise Not_found
+
+  let find_last_opt f s =
+    fold
+      (fun elt acc ->
+        match acc with
+        | None -> if f elt then Some elt else None
+        | Some witness ->
+            if witness.tag >= elt.tag then acc else
+            if f elt then Some elt else acc)
+      s None
+
+  let find_last f s =
+    match find_last_opt f s with
+      | Some elt -> elt
+      | None -> raise Not_found
+
   (*s Another nice property of Patricia trees is to be independent of the
       order of insertion. As a consequence, two Patricia trees have the
       same elements if and only if they are structurally equal.
@@ -959,4 +1039,26 @@ module Hset = struct
         intersect s1 (if zero_bit p1 m2 then l2 else r2)
       else
         false
+
+  let disjoint s1 s2 = not (intersect s1 s2)
+
+  let find_any (type a) f (s : a t) =
+    let exception Found of a elt in
+    try
+      iter (fun elt -> if f elt then raise (Found elt)) s;
+      raise Not_found
+    with Found elt -> elt
+  let find_any_opt (type a) f (s : a t) =
+    let exception Found of a elt in
+    try
+      iter (fun elt -> if f elt then raise (Found elt)) s;
+      None
+    with Found elt -> Some elt
+
+  let bind f s = fold (fun elt s -> union (f elt) s) s empty
+
+  let is_singleton = function
+    | Leaf elt -> Some elt
+    | _ -> None
+
 end
diff --git a/hashcons.mli b/hashcons.mli
@@ -169,27 +169,60 @@ module Hset : sig
   val diff : 'a t -> 'a t -> 'a t
   val equal : 'a t -> 'a t -> bool
   val compare : 'a t -> 'a t -> int
-  val elements : 'a t -> 'a elt list
   val choose : 'a t -> 'a elt
+  val choose_opt : 'a t -> 'a elt option
   val cardinal : 'a t -> int
-  val iter : ('a elt -> unit) -> 'a t -> unit
-  val fold : ('a elt -> 'b -> 'b) -> 'a t -> 'b -> 'b
   val for_all : ('a elt -> bool) -> 'a t -> bool
   val exists : ('a elt -> bool) -> 'a t -> bool
-  val filter : ('a elt -> bool) -> 'a t -> 'a t
   val partition : ('a elt -> bool) -> 'a t -> 'a t * 'a t
+  val disjoint : 'a t -> 'a t -> bool
+  val find : 'a elt -> 'a t -> 'a elt
+  val find_opt :  'a elt -> 'a t -> 'a elt option
+  val add_seq : 'a elt Seq.t -> 'a t -> 'a t
+  val of_seq : 'a elt Seq.t -> 'a t
+  val of_list : 'a elt list -> 'a t
+  val split : 'a elt -> 'a t -> 'a t * bool * 'a t
+
+  (*s Warning: [iter], [fold], [map], [filter] and [map_filter] do NOT iterate
+      over element order. Similarly, [elements] and [to_seq] are not sorted. *)
+  val iter : ('a elt -> unit) -> 'a t -> unit
+  val fold : ('a elt -> 'b -> 'b) -> 'a t -> 'b -> 'b
+  val map : ('a elt -> 'b elt) -> 'a t -> 'b t
+  val filter : ('a elt -> bool) -> 'a t -> 'a t
+  val filter_map : ('a elt -> 'b elt option) -> 'a t -> 'b t
+  val elements : 'a t -> 'a elt list
+  val to_seq : 'a t -> 'a elt Seq.t
 
-  (*s Warning: [min_elt] and [max_elt] are linear w.r.t. the size of the
-      set. In other words, [min_elt t] is barely more efficient than [fold
-      min t (choose t)]. *)
+  (*s Warning: [min_elt], [max_elt] and the [_opt] versions are linear w.r.t.
+      the size of the set. In other words, [min_elt t] is barely more efficient
+      than [fold min t (choose t)]. *)
   val min_elt : 'a t -> 'a elt
+  val min_elt_opt : 'a t -> 'a elt option
   val max_elt : 'a t -> 'a elt
+  val max_elt_opt : 'a t -> 'a elt option
+
+  (*s [find_first], [find_last] are linear time and can call [f] an arbitrary
+      number of times, and not necessarily on elements smaller/larger
+      than the witness. *)
+  val find_first : ('a elt -> bool) -> 'a t -> 'a elt
+  val find_first_opt : ('a elt -> bool) -> 'a t -> 'a elt option
+  val find_last : ('a elt -> bool) -> 'a t -> 'a elt
+  val find_last_opt : ('a elt -> bool) -> 'a t -> 'a elt option
 
   (*s Additional functions not appearing in the signature [Set.S] from ocaml
       standard library. *)
 
   (* [intersect u v] determines if sets [u] and [v] have a non-empty
      intersection. *)
   val intersect : 'a t -> 'a t -> bool
+
+  (* Faster finds when order doesn't matter *)
+  val find_any : ('a elt -> bool) -> 'a t -> 'a elt
+  val find_any_opt : ('a elt -> bool) -> 'a t -> 'a elt option
+
+  val is_singleton : 'a t -> 'a elt option
+  (* Check if the set is a singleton, if so return unique element *)
+
+  val bind : ('a elt -> 'b t) -> 'a t -> 'b t
 end
 
