let edges = [
  (1, 0, 1); (2, 1, 2); (3, 0, 3); (4, 2, 3); (5, 3, 4); (6, 0, 4)
]

let make_uf n = Array.init n (fun i -> i)

let rec find p x =
  if p.(x) = x then x
  else begin
    let r = find p p.(x) in
    p.(x) <- r;
    r
  end

let union p x y =
  let rx = find p x in
  let ry = find p y in
  if rx <> ry then begin p.(rx) <- ry; true end
  else false

let mst_weight n es =
  let sorted = List.sort (fun (w1, _, _) (w2, _, _) -> compare w1 w2) es in
  let p = make_uf n in
  let total = ref 0 in
  List.iter (fun (w, u, v) ->
    if union p u v then total := !total + w
  ) sorted;
  !total

;;

mst_weight 5 edges
