Skip to content
Permalink
master
Switch branches/tags
Go to file
 
 
Cannot retrieve contributors at this time
module UnionFind = struct
type t =
{ parent: int array;
rank: int array;
size: int array;
num_sets: int ref;
}
let create n =
{ parent = Array.init n (fun x -> x);
rank = Array.make n 0;
size = Array.make n 1;
num_sets = ref n;
}
let find { parent } u =
let rec aux u =
if parent.(u) = u then u
else begin
parent.(u) <- aux parent.(u);
parent.(u)
end
in aux u
let is_same_set dsu u v = find dsu u = find dsu v
let union dsu u v =
let link x p =
dsu.parent.(x) <- p;
dsu.size.(p) <- dsu.size.(p) + dsu.size.(x);
if dsu.rank.(x) = dsu.rank.(p) then
dsu.rank.(p) <- dsu.rank.(p) + 1
in
let pu, pv = find dsu u, find dsu v in
if pu = pv then ()
else begin
decr dsu.num_sets;
if dsu.rank.(pu) > dsu.rank.(pv) then link pv pu
else link pu pv
end
let num_disjoint_sets { num_sets } = !num_sets
let size_of_set dsu u = dsu.size.(find dsu u)
end
(** [kruskal n edges] returns the total weight of a MST of the input graph [(0..n-1, edges)] *)
let kruskal n edges =
let dsu = UnionFind.create n in
let cost = ref 0 in
edges
|> List.sort (fun (_, _, w1) (_, _, w2) -> compare w1 w2)
|> List.iter (fun (u, v, w) ->
if not (UnionFind.is_same_set dsu u v) then begin
UnionFind.union dsu u v;
cost := !cost + w
end
);
!cost
open Scanf
open Printf
let sc = object method nf = bscanf Scanning.stdin end
let () =
let (n, m) = sc#nf "%d %d\n" (fun x y -> (x, y)) in
let edges = List.init m (fun _ -> sc#nf "%d %d %d\n" (fun u v w -> (u, v, w))) in
kruskal n edges |> printf "MST cost = %d (Kruskal's)\n"