randomized order statistic selection in OCaml

rselect.ml

let swap arr i j =
  let tmp = arr.(j) in
  arr.(j) <- arr.(i);
  arr.(i) <- tmp

let partition arr l h =
  let pivot = arr.(h) in
  let i = ref @@ (l - 1) in
  for j = l to h - 1 do
    if arr.(j) <= pivot then begin
      incr i;
      swap arr !i j
    end
  done;
  incr i;
  swap arr !i h;
  !i

let partition_random arr l h =
  let pivot_index = l + Random.int (h - l) in
  swap arr pivot_index h;
  partition arr l h

let rselect i xs =
  if i < 1 then invalid_arg "'i' must be positive.";
  let rec rselect' arr l h i =
    if i < 0 then invalid_arg "'i' must be non-negative.";
    if l = h then arr.(l)
    else if l < h then
      let pivot_index = partition_random arr l h in
      let k = pivot_index - l in
      if i = k then arr.(pivot_index)
      else if i < k then rselect' arr l (pivot_index - 1) i
      else rselect' arr (pivot_index + 1) h (i - k - 1)
    else invalid_arg "Something went wrong."
  in
  let arr = Array.of_list xs in
  rselect' arr 0 (List.length xs - 1) (i - 1)

let shuffle xs =
  List.map (fun x -> (Random.bits (), x)) xs
  |> List.sort compare |> List.map snd

let () =
  Random.self_init ();
  let n = 10 in
  for i = 1 to n do
    List.init n (fun x -> x + 1) |> shuffle |> rselect i |> Printf.printf "%d";
    print_newline ()
  done