(*
  Union-find data structure with the path compression optimization.
*)
open Array

type node = { mutable parent : int; 
	      mutable rank : int; }
type union_find = node array


(* Add a new set to the data structure *)
let make_set (uf : union_find) : int * union_find =
  let n = length uf in
    (n, append uf (of_list [{parent = n; rank = 0}]))

(* Find the representative for the set containing v.
   Also, this performs path compression, mutating the union-find
   data structure.
*)
let find (v : int) (uf : union_find) : int =
  let old = v in
  let rec loop ancestor v =
    (if (ancestor = v) then
      let rec inner_loop old v =
	if (ancestor = v) then
	  ancestor
	else (
	  uf.(old).parent <- ancestor;
	  inner_loop v uf.(v).parent)  in
      let v = uf.(old).parent in
	inner_loop old v
     else
       loop uf.(ancestor).parent ancestor) in
  let ancestor = uf.(v).parent in
    loop ancestor v
    
let union (x : int) (y : int) (uf : union_find) =
  let x = find x uf and y = find y uf in
    (if (x = y) then ()
     else if (uf.(x).rank > uf.(y).rank) then
       uf.(y).parent <- x
     else (
       uf.(x).parent <- y;
       if (uf.(x).rank = uf.(y).rank) then
	 uf.(y).rank <- uf.(y).rank + 1
       else
	 ()))

(* Use the first argument as the representative *)
let union_fst_rep (x : int) (y : int) (uf : union_find) =
  let x = find x uf and y = find y uf in
    (if (x = y) then 
       ()
     else
       uf.(y).parent <- x)