A permutation generator with Johnson Trotter algorithm

johnson_permutation_generator.ml

type direction = L | R

let attach_direction a = Array.map (fun x -> (x, L)) a

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

let is_movable a i =
  let x,d = a.(i) in
  match d with
    | L -> if i > 0 && x > (fst a.(i-1)) then true else false
    | R -> if i < Array.length a - 1 && x > (fst a.(i+1)) then true else false
    
let move a i =
  let x,d = a.(i) in
  if is_movable a i then 
    match d with
      | L -> swap a i (i-1)
      | R -> swap a i (i+1)
  else
    failwith "not movable"
    
let scan_movable_largest a =
  let rec aux acc i =
    if i >= Array.length a then acc
    else if not (is_movable a i) then aux acc (i+1)
    else
      let x,_ = a.(i) in
      match acc with
        | None -> aux (Some i) (i+1)
        | Some j -> aux (if x < fst(a.(j)) then acc else Some i) (i+1)
  in
  aux None 0
  
let flip = function | L -> R | R -> L

let scan_flip_larger x a = 
  Array.iteri (fun i (y, d) -> if y > x then a.(i) <- y,flip d) a
  
let permutations_generator l =
  let a = Array.of_list l |> attach_direction in
  let r = ref (Some l) in
  let next () = 
    let p = !r in
    (match scan_movable_largest a with
      | None -> r := None
      | Some i -> 
        let x, _ = a.(i) in (
        move a i; 
        scan_flip_larger x a;
        r := Some (Array.map fst a |> Array.to_list)));
    p
  in
  next

(* example *)
let generator = permutations_generator [1;2;3]