Solving 3-In-A-Row using Z3

3-In-A-Row puzzle

Can you fill the grid with Blue and White squares without creating a 3-In-A-Row of the same colour? Each row and column has an equal number of Blue and White squares. Blue is represented by “X” and white is represented by “O” in the picture below

Solution using Z3

from z3 import *

class ThreeInARowSolver:
    def __init__(self, n, puzzle):
        self.n = n
        self.input = puzzle
        self.X = [[Int("x_%s_%s" % (i, j)) for j in range(n)]
                      for i in range(n)]
        self.s = Solver()
        self.s.add([And(0 <= self.X[i][j], self.X[i][j]<= 1) for i in range(n) for j in range(n)])

    def set_constraints(self):
        for i,j,v in self.input:
            self.s.add(And(self.X[i][j] == v))
        for i in range(self.n):
            self.s.add(And(sum(self.X[i]) == self.n/2))
        for j in range(self.n):
            self.s.add(And(sum([self.X[i][j] for i in range(self.n)]) == self.n/2))
        for i in range(0, self.n):
            for j in range(0, self.n-2):
                self.s.add(And(self.X[i][j] + self.X[i][j+1] + self.X[i][j+2] != 0))
                self.s.add(And(self.X[i][j] + self.X[i][j+1] + self.X[i][j+2] != 3))
        for j in range(0, self.n):
            for i in range(0, self.n-2):
                self.s.add(And(self.X[i][j] + self.X[i+1][j] + self.X[i+2][j] != 0))
                self.s.add(And(self.X[i][j] + self.X[i+1][j] + self.X[i+2][j] != 3))

    def output_solution(self):
        m = self.s.model()
        for i in range(self.n):
            print(" ".join(["X" if m.evaluate(self.X[i][j])==1 else "O" for j in range(self.n)]))

    def solve(self):
        self.set_constraints()
        if self.s.check() == sat:
            self.output_solution()
        else:
            print(self.s)
            print("Failed to solve.")

""" The puzzle is encoded using a 1 for an X and an O for an O. The positions of the X's and O's are captured as a list of 3-tuples (row, column, symbol)."""

puzzle = [
    (0,4,1),
    (1,3,1),
    (3,5,1),
    (0,1,0),
    (2,5,0),
    (4,0,0),
    (5,0,0),
    (5,1,0)
]

ts = ThreeInARowSolver(6, puzzle)
ts.solve()

Solution

X O X O X O
O X O X O X
X X O O X O
X O X O O X
O X O X X O
O O X X O X