(* 使い方 
  ocamlopt.opt str.cmxa kruskal.ml -o kruskal
  ./kruskal 1000 2000 "graph2.txt" 
 一つめの引数は頂点数、二つ目は辺の数、３つめはデータファイル。
  データファイルはランダムグラフサーバの形式にそっているが最後の
  空白行を消さないとエラーを出す 。
  出力結果はかかった時間とデータファイルの頂点数と最小全域木の頂点数。 
*)



(* kruskal algorithm *)

type heapelement = { from:int; into:int; value:float};;

let rec findgroup k parent =
  if k=parent.(k) then
    k
  else 
    let a = findgroup parent.(k) parent in
      parent.(k) <- a;
      a;;


let rec simple_findgroup k parent =
  if k<>parent.(k) then
    findgroup parent.(k) parent
  else 
    k;;
  

let uniongroup a b parent treesize =
  let p = findgroup a parent in
  let q = findgroup b parent in
    if treesize.(p)>=treesize.(q) then
      begin
	parent.(q) <- p;
	treesize.(p) <- treesize.(p) + treesize.(q) 
      end
    else
      begin
	parent.(p) <- q;
	treesize.(q) <- treesize.(p) + treesize.(q)
      end;;

(* read data from file and heap the data *)
let init_element element m filename =
  let a = open_in filename in
    for i=1 to m do
      let strlist = Str.split_delim (Str.regexp " ") (input_line a)  in
      let vertex1 = (int_of_string (List.nth strlist 0)) in
      let vertex2 = (int_of_string (List.nth strlist 1)) in
      let weight = (float_of_string (List.nth strlist 2)) in
      let s = ref i in
      let p = ref (!s/2) in  (* parent *)
	(* set value on last heap *)
	element.(i) <- {from=vertex1;into=vertex2;value=weight};
	(* shift up *)
	while (!s>=2 && element.(!p).value > element.(!s).value) do
	  let w = element.(!p) in (* ? *)
	    element.(!p) <- element.(!s);
	    element.(!s) <- w;
	    s := !p;
	    p := !s/2;
	done 
    done;; 

(* n:vertex number m: edge number *)
let kruskal n m filename =
  let element = Array.make (m+1) {from=0;into=0;value=0.} in 
  let parent = Array.make (n+1) 1 in
  let treesize = Array.make (n+1) 1 in
  let return_value = ref [] in
    (* init parent *)
    for i=1 to n do
      parent.(i) <- i
    done;
    (* init element *)
    init_element element m filename;
    for i=1 to m do
      let u = element.(i).from in
      let v = element.(i).into in
	if (findgroup u parent)<>(findgroup v parent) then
	  begin
	    uniongroup u v parent treesize;
	    return_value := element.(i)::!return_value;
	  end;
    done;
    !return_value;;

(*
let test n m filename =
  let a = Sys.time () in
  let b = ref 0.0 in
    ignore(kruskal n m filename);
    b := Sys.time();
    !b -. a;;
*)


let check1 name arg1 arg2 =
  let a = open_in name in
  let b = Array.make arg1 0 in
    for i=1 to arg2 do
      let strlist = Str.split_delim (Str.regexp " ") (input_line a)  in
      let vertex1 = (int_of_string (List.nth strlist 0)) in
      let vertex2 = (int_of_string (List.nth strlist 1)) in
	b.(vertex1-1) <- 1;
	b.(vertex2-1) <- 1;
    done;
    Array.fold_left (fun a b -> a + b) 0 b;;


let check2 result arg1=
  let rec sub_fun1 element array =
    match element with
	hd::tl -> 
	  array.(hd.from-1) <- 1;
	  array.(hd.into-1) <- 1; 
	  sub_fun1 tl array;
      | [] ->  () in
  let a = Array.make arg1 0 in
    sub_fun1 result a;
    Array.fold_left (fun a b -> a + b) 0 a;;

let _  =
  let a = Sys.time () in
  let b = ref 0.0 in
  let arg1 = int_of_string Sys.argv.(1) in
  let arg2 = int_of_string Sys.argv.(2) in
  let str = Sys.argv.(3) in
  let result = kruskal arg1 arg2 str in
    b := Sys.time();
    print_string "Time:";
    print_float (!b -. a);
    print_newline ();
    print_string "data file check:";
    print_int (check1 str arg1 arg2);
    print_newline ();
    print_string "result check:";
    print_int (check2 result arg1);
    print_newline ();;