practice of https://github.com/termite-analyser/z3overlay

sudokusolver.ml

module Z3 = ZZ3.Make (struct let ctx = Z3.mk_context [] end)

open Z3

let solve board =

  let cells =
    Array.init 9 (fun i ->
      Array.init 9 (fun j ->
        Symbol.declare Int (Printf.sprintf "%d%d" i j)))
  in

  (* Return following cells as a list

     (i, j),          (i, j + 1),          ..., (i, j + sx - 1),
     (i + 1, j),      (i + 1, j + 1),      ..., (i + 1, j + sx - 1),
     ...
     (i + sy - 1, j), (i + sy - 1, j + 1), ..., (i + sy - 1, j + sx - 1)

     where
       sy = size_y
       sx = size_x
  *)
  let extract ~size_y ~size_x i j =
    let rec loop acc y x =
      if y = size_y
      then acc
      else loop
             (cells.(i + y).(j + x) :: acc)
             (y + (x + 1) / size_x)
             ((x + 1) mod size_x)
    in
    loop [] 0 0
  in

  let solver = Solver.make () in

  (* all numbers in cells are 1 <= and <= 9 *)
  for i = 0 to 8 do
    for j = 0 to 8 do
      let c = T.(int 1 <= !(cells.(i).(j)) && !(cells.(i).(j)) <= int 9) in
      Solver.add ~solver c
    done
  done;

  (* constraints on each row *)
  for i = 0 to 8 do
    let c = extract ~size_y:1 ~size_x:9 i 0
            |> List.map T.symbol
            |> T.distinct
    in
    Solver.add ~solver c
  done;

  (* constraints on each column *)
  for j = 0 to 8 do
    let c = extract ~size_y:9 ~size_x:1 0 j
            |> List.map T.symbol
            |> T.distinct
    in
    Solver.add ~solver c
  done;

  (* constraints on each block *)
  for i = 0 to 2 do
    for j = 0 to 2 do
      let c = extract ~size_y:3 ~size_x:3 (i * 3) (j * 3)
              |> List.map T.symbol
              |> T.distinct
      in
      Solver.add ~solver c
    done
  done;

  (* assign initial numbers as constraints *)
  List.iteri (fun i row ->
    List.iteri (fun j num ->
      if num <> " "
      then
        let c = T.(!(cells.(i).(j)) = int (int_of_string num)) in
        Solver.add ~solver c
    ) row
  ) board;

  let model = match (Solver.check ~solver []) with
    | Unsat _ | Unkown _ -> failwith "no solution"
    | Sat (lazy model) -> model
  in

  for i = 0 to 8 do
    for j = 0 to 8 do
      print_int (Model.get_value ~model cells.(i).(j) |> Z.to_int)
    done;
    print_newline ();
  done

let () =
  let board = [ "    1    "
              ; "5      2 "
              ; " 4 8  6  "
              ; "       35"
              ; " 689     "
              ; "  7      "
              ; "    3    "
              ; "    52   "
              ; " 9    7  "
              ]
              |> List.map (Str.split (Str.regexp ""))
  in
  solve board