Open Source Adventures: Episode 70: Crystal Z3 Solver for Nonograms Puzzle

4 2 1 2 2
4 1 1 3
4 2 2 3
8 10
6 3 9
3 1 1 9
2 1 1 8
2 1 1 2 2 5
1 1 4 1 3
2 9 6
13 7
4 1 6 4 2
1 1 3 3 3
1 3 1 1 2
2 8 2
5 1 4 2
5 3 1 2
5 2 4
6 2 3
5 1 2
3 4 2
4 2 3 2
6 3 4 3
8 4 1
1 1 4 4
-
6 1 2
1 5 1 3
6 8 3
5 5 10
5 4 9
4 2 9
1 1 5 1
1 3 1 2 2
4 2 1 2
1 1 6 2
1 3 8
10
5 5 3
9 6
4 5 1 5
6 3 2 5
4 1
8
1 6 2 1
4 5 1
5 5 1 1
7 9 1
11 2 2 3
4 6 6
5 6

#!/usr/bin/env crystal

require "z3"

class Nonograms
  @col_hints : Array(Array(Int32))
  @row_hints : Array(Array(Int32))

  def initialize(path)
    col, row = File.read(path).split("-\n", 2)
    @col_hints = col.chomp.split("\n").map{|l| l.split.map(&.to_i)}
    @row_hints = row.chomp.split("\n").map{|l| l.split.map(&.to_i)}
    @xsize = @col_hints.size
    @ysize = @row_hints.size
    @solver = Z3::Solver.new
    @cells = {} of Tuple(Int32, Int32) => Z3::BoolExpr
  end

  def column_cells(x)
    (0...@ysize).map{|y| @cells[{x,y}]}
  end

  def row_cells(y)
    (0...@xsize).map{|x| @cells[{x,y}]}
  end

  def setup_groups(name, hints, cells)
    count = hints.size
    starts = (0...count).map{|i| Z3.int("#{name} start #{i}") }
    ends = (0...count).map{|i| Z3.int("#{name} end #{i}") }

    # Every group is on the board
    count.times do |i|
      @solver.assert starts[i] >= 0
      @solver.assert ends[i] <= cells.size
    end

    # Every group has correct size
    hints.each_with_index do |hint, i|
      @solver.assert starts[i] + hint == ends[i]
    end

    # There is a gap between groups
    (count-1).times do |i|
      @solver.assert starts[i+1] >= ends[i] + 1
    end

    # Each filled cell belongs to one of the groups
    cells.each_with_index do |c, j|
      cell_j_in_group_i = count.times.map{|i| (j >= starts[i]) & (j < ends[i]) }
      # this is fairly awkward and we should probably add Z3.or()
      cell_j_in_any_group = count == 0 ? false : cell_j_in_group_i.reduce{|a,b| a|b}
      @solver.assert c == cell_j_in_any_group
    end
  end

  def call
    # Setup cell variables
    @xsize.times do |x|
      @ysize.times do |y|
        @cells[{x,y}] = Z3.bool("cell #{x},#{y}")
      end
    end

    # Setup groups
    @col_hints.each_with_index do |hints, x|
      setup_groups("c#{x}", hints, column_cells(x))
    end

    @row_hints.each_with_index do |hints, y|
      setup_groups("r#{y}", hints, row_cells(y))
    end

    # Print the solution
    if @solver.satisfiable?
      model = @solver.model
      @ysize.times do |y|
        @xsize.times do |x|
          v = (model[@cells[{x, y}]].to_s == "true")
          if v
            print "#"
          else
            print "."
          end
        end
        print "\n"
      end
    else
      puts "No solution"
    end
  end
end

Nonograms.new(ARGV[0]).call

$ ./nonograms 1.txt
...######......#......##.
#..#####.......#......###
######.....########..###.
#####..#####..##########.
#####....####..#########.
.####.....##....#########
...#......#....#####...#.
...#.....###..#...##..##.
....####.##...#......##..
#...#....######......##..
#...###..########........
.........##########......
.......#####..#####...###
.......#########...######
..####..#####.#.....#####
######..###.##......#####
...####.............#....
########.................
#.######..##....#........
...####.#####...........#
...#####.#####..#.......#
.#######.#########......#
###########...##.##...###
####..######.....######..
.......#####.....######..