aboutsummaryrefslogtreecommitdiffhomepage
path: root/lib/unionfind.ml
blob: 6e131d8fbd6f1b33c77200aa6cb8904a19573484 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
(************************************************************************)
(*  v      *   The Coq Proof Assistant  /  The Coq Development Team     *)
(* <O___,, *   INRIA - CNRS - LIX - LRI - PPS - Copyright 1999-2016     *)
(*   \VV/  **************************************************************)
(*    //   *      This file is distributed under the terms of the       *)
(*         *       GNU Lesser General Public License Version 2.1        *)
(************************************************************************)

(** An imperative implementation of partitions via Union-Find *)

(** Paths are compressed imperatively at each lookup of a
    canonical representative. Each union also modifies in-place
    the partition structure.

    Nota: For the moment we use Pervasive's comparison for
    choosing the smallest object as representative. This could
    be made more generic.
*)



module type PartitionSig = sig

  (** The type of elements in the partition *)
  type elt

  (** A set structure over elements *)
  type set

  (** The type of partitions *)
  type t

  (** Initialise an empty partition *)
  val create : unit -> t

  (** Add (in place) an element in the partition, or do nothing
      if the element is already in the partition. *)
  val add : elt -> t -> unit

  (** Find the canonical representative of an element.
      Raise [not_found] if the element isn't known yet. *)
  val find : elt -> t -> elt

  (** Merge (in place) the equivalence classes of two elements.
      This will add the elements in the partition if necessary. *)
  val union : elt -> elt -> t -> unit

  (** Merge (in place) the equivalence classes of many elements. *)
  val union_set : set -> t -> unit

  (** Listing the different components of the partition *)
  val partition : t -> set list

end

module type SetS =
sig
  type t
  type elt
  val singleton : elt -> t
  val union : t -> t -> t
  val choose : t -> elt
  val iter : (elt -> unit) -> t -> unit
end

module type MapS =
sig
  type key
  type +'a t
  val empty : 'a t
  val find : key -> 'a t -> 'a
  val add : key -> 'a -> 'a t -> 'a t
  val mem : key -> 'a t -> bool
  val fold : (key -> 'a -> 'b -> 'b) -> 'a t -> 'b -> 'b
end

module Make (S:SetS)(M:MapS with type key = S.elt) = struct

  type elt = S.elt
  type set = S.t

  type node =
    | Canon of set
    | Equiv of elt

  type t = node ref M.t ref

  let create () = ref (M.empty : node ref M.t)

  let fresh x p =
    let node = ref (Canon (S.singleton x)) in
    p := M.add x node !p;
    x, node

  let rec lookup x p =
    let node = M.find x !p in
    match !node with
      | Canon _ -> x, node
      | Equiv y ->
	let ((z,_) as res) = lookup y p in
	if not (z == y) then node := Equiv z;
	res

  let add x p = if not (M.mem x !p) then ignore (fresh x p)

  let find x p = fst (lookup x p)

  let canonical x p = try lookup x p with Not_found -> fresh x p

  let union x y p =
    let ((x,_) as xcan) = canonical x p in
    let ((y,_) as ycan) = canonical y p in
    if x = y then ()
    else
      let xcan, ycan = if x < y then xcan, ycan else ycan, xcan in
      let x,xnode = xcan and y,ynode = ycan in
      match !xnode, !ynode with
	| Canon lx, Canon ly ->
	  xnode := Canon (S.union lx ly);
	  ynode := Equiv x;
	| _ -> assert false

  let union_set s p =
    try
      let x = S.choose s in
      S.iter (fun y -> union x y p) s
    with Not_found -> ()

  let partition p =
    List.rev (M.fold
		(fun x node acc -> match !node with
		  | Equiv _ -> acc
		  | Canon lx -> lx::acc)
		!p [])

end