(* persistent union-find = Tarjan's algorithm with persistent arrays *)

module Make(A : Sig.PolyPersistentArray) : Sig.UnionFind = struct

  type t = { 
    rank: int A.t;
    mutable father: int A.t;
  }
      
  let create n = 
    { rank = A.create n 0;
      father = A.init n (fun i -> i) }

  let rec find_aux f i = 
    let fi = A.get f i in
    if fi == i then 
      f, i
    else 
      let f, r = find_aux f fi in 
      let f = A.set f i r in
      f, r
      
  let find h x = 
    let f,rx = find_aux h.father x in h.father <- f; rx
  
  let union h x y = 
    let rx = find h x in
    let ry = find h y in
    if rx != ry then begin
      let rxc = A.get h.rank rx in
      let ryc = A.get h.rank ry in
      if rxc > ryc then
	{ h with father = A.set h.father ry rx }
      else if rxc < ryc then
	{ h with father = A.set h.father rx ry }
      else
	{ rank = A.set h.rank rx (rxc + 1);
	  father = A.set h.father ry rx }
    end else
      h

end