Skip to content
Switch branches/tags
Go to file
Cannot retrieve contributors at this time
module PQ = Set.Make (struct
type t = int * int
let compare = compare
module IntS = Set.Make (struct
type t = int
let compare = compare
(** [prim n edges] returns the total weight of a MST of the input graph [(0..n-1, edges)] *)
let prim n edges =
let adj = Array.make n [] in
edges |> List.iter (fun (u, v, w) ->
adj.(u) <- (w, v) :: adj.(u);
adj.(v) <- (w, u) :: adj.(v)
let rec loop taken candidates cost =
match PQ.min_elt_opt candidates with
| None -> cost
| Some ((w, u) as e) ->
let candidates = PQ.remove e candidates in
if IntS.mem u taken then loop taken candidates cost
else loop (IntS.add u taken)
(candidates |> PQ.add_seq (List.to_seq adj.(u)))
(cost + w)
in loop (IntS.singleton 0) (PQ.of_list adj.(0)) 0
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
prim n edges |> printf "MST cost = %d (Prim's)\n"