Union-Find data-structure in OCaml

union_find.ml

type 'a member =
  { mutable parent : 'a member
  ; mutable rank : int
  ; value : 'a
  }

let make_set value =
  let rec m = { parent = m; rank = 0; value } in
  m
;;

let rec find_set member =
  if member.parent != member
  then (* path compression *)
    member.parent <- find_set member.parent;
  member.parent
;;

let link x y =
  (* rank heuristic *)
  if x.rank > y.rank
  then y.parent <- x
  else (
    x.parent <- y;
    if x.rank == y.rank then y.rank <- y.rank + 1)
;;

let union x y = link (find_set x) (find_set y)
let in_same_set x y = find_set x == find_set y
let is_representative member = member.parent == member

union_find.mli

(* Union-Find. See "Chapter 21. Data Structures for Disjoint Sets" in [1].

   [1]: "Introduction to Algorithms". 3rd Ed. Cormen et al. *)

type 'a member

(* Create a new set that has the single element as representative. *)
val make_set : 'a -> 'a member

(* Find the member that represents the set of the passed member. *)
val find_set : 'a member -> 'a member

(* Merges the two sets that contain the passed members into a union of both
   sets. *)
val union : 'a member -> 'a member -> unit

(* Are the two elements part of the same set. *)
val in_same_set : 'a member -> 'a member -> bool

(* Is the member that represents the set in the Union-Find? *)
val is_representative : 'a member -> bool

union_find_test.ml

open Base
open Stdio

let print_set_forest sets =
  let n = Array.length sets in
  for i = 0 to n - 1 do
    let x = sets.(i) in
    for j = 0 to n - 1 do
      let y = sets.(j) in
      let same_set = Union_find.in_same_set x y in
      if same_set then printf "%d %d\n" i j
    done
  done
;;

let make_sets n = Array.init n ~f:(fun i -> Union_find.make_set i)

let%expect_test "union-find" =
  let sets = make_sets 3 in
  print_set_forest sets;
  [%expect {|
  0 0
  1 1
  2 2
  |}];
  Union_find.union sets.(0) sets.(1);
  print_set_forest sets;
  [%expect {|
  0 0
  0 1
  1 0
  1 1
  2 2
  |}];
  Union_find.union sets.(0) sets.(2);
  print_set_forest sets;
  [%expect {|
  0 0
  0 1
  0 2
  1 0
  1 1
  1 2
  2 0
  2 1
  2 2
  |}];
  let sets2 = make_sets 3 in
  Union_find.union sets2.(0) sets2.(1);
  Union_find.union sets2.(1) sets2.(2);
  print_set_forest sets2;
  [%expect {|
  0 0
  0 1
  0 2
  1 0
  1 1
  1 2
  2 0
  2 1
  2 2
  |}]
;;

let%expect_test "union-all" =
  let sets = make_sets 10000 in
  let outcome i j =
    if Union_find.in_same_set sets.(i) sets.(j) then "joint" else "disjoint"
  in
  for i = 1 to 1 do
    let j = i - 1 in
    print_endline (outcome j i);
    [%expect {| disjoint |}];
    Union_find.union sets.(j) sets.(i);
    print_endline (outcome j i);
    [%expect {| joint |}]
  done
;;