ligo/src/union_find/Partition3.ml
2019-09-27 14:55:09 +02:00

87 lines
2.8 KiB
OCaml

(* Destructive implementation of union/find with height-balanced
forests but without path compression: O(n*log(n)). *)
module Make (Item: Partition.Item) =
struct
type item = Item.t
type repr = item (** Class representatives *)
let equal i j = Item.compare i j = 0
type height = int
(** Each equivalence class is implemented by a Catalan tree linked
upwardly and otherwise is a link to another node. Those trees
are height-balanced. The type [node] implements nodes in those
trees. *)
type node = {item: item; mutable height: int; mutable parent: node}
module ItemMap = Map.Make (Item)
(** The type [partition] implements a partition of classes of
equivalent items by means of a map from items to nodes of type
[node] in trees. *)
type partition = node ItemMap.t
type t = partition
let empty = ItemMap.empty
(** The function [repr] is faster than a persistent implementation
in the worst case because, in the latter case, the cost is O(log n)
for accessing each node in the path to the root, whereas, in the
former, only the access to the first node in the path incurs a cost
of O(log n) -- the other nodes are accessed in constant time by
following the [next] field of type [node]. *)
let seek (i: item) (p: partition) : node =
let rec find_root node =
if node.parent == node then node else find_root node.parent
in find_root (ItemMap.find i p)
let repr item partition = (seek item partition).item
let is_equiv (i: item) (j: item) (p: partition) =
equal (repr i p) (repr j p)
let get_or_set item (p: partition) =
try seek item p, p with
Not_found -> let rec loop = {item; height=0; parent=loop}
in loop, ItemMap.add item loop p
let link src dst = src.parent <- dst
let equiv (i: item) (j: item) (p: partition) : partition =
let ni,p = get_or_set i p in
let nj,p = get_or_set j p in
let hi,hj = ni.height, nj.height in
let () =
if not (equal ni.item nj.item)
then if hi > hj
then link nj ni
else (link ni nj; nj.height <- max hj (hi+1))
in p
let alias (i: item) (j: item) (p: partition) : partition =
let ni,p = get_or_set i p in
let nj,p = get_or_set j p in
let hi,hj = ni.height, nj.height in
let () =
if not (equal ni.item nj.item)
then if hi = hj || equal ni.item i
then (link ni nj; nj.height <- max hj (hi+1))
else if hi < hj then link ni nj
else link nj ni
in p
(* Printing *)
let print p =
let print _ node =
Printf.printf "%s,%d -> %s,%d\n"
(Item.to_string node.item) node.height
(Item.to_string node.parent.item) node.parent.height
in ItemMap.iter print p
end