Sudoku CNF

sudoku.py

import collections
import itertools as it
import subprocess
import sys

class Sudoku:
	def __init__(self, size):
		self.width = size*size
		self.height = size*size
		self._opts = size*size
		self.big_squares = (size, size, size)
		self.constraints = []

	def rows(self):
		return range(self.width)

	def cols(self):
		return range(self.height)

	def squares(self):
		return [(x, y) for y in self.cols() for x in self.rows()]

	def opts(self):
		return [o for o in range(self._opts)]

	def var(self, x, y, o):
		return y*self.width*self._opts + x*self._opts + o + 1

	def decode_var(self, v):
		v -= 1
		o = v % self._opts
		v //= self._opts
		x = v % self.width
		v //= self.width
		y = v
		return (x, y, o)

	def add_base_constraints(self):
		# Squares can't have two values
		for x, y in self.squares():
			for o1, o2 in it.combinations(self.opts(), 2):
				self.add_constraint([-self.var(x, y, o1), -self.var(x, y, o2)])

		# Rows can't have duplicate values
		for x1, x2 in it.combinations(self.rows(), 2):
			for y in self.cols():
				for o in self.opts():
					self.add_constraint([-self.var(x1, y, o), -self.var(x2, y, o)])

		# Columns can't have duplicate values
		for x in self.rows():
			for y1, y2 in it.combinations(self.cols(), 2):
				for o in self.opts():
					self.add_constraint([-self.var(x, y1, o), -self.var(x, y2, o)])

		# Big squares can't have duplicate values
		if self.big_squares is not None:
			brows, bcols, bsize = self.big_squares
			for bx, by in it.product(range(brows), range(bcols)):
				for o in self.opts():
					for ((sx1, sy1), (sx2, sy2)) in it.combinations(it.product(range(bsize), range(bsize)), 2):
						x1 = bx*bsize + sx1
						y1 = by*bsize + sy1
						x2 = bx*bsize + sx2
						y2 = by*bsize + sy2
						self.add_constraint([-self.var(x1, y1, o), -self.var(x2, y2, o)])

		# Squares must have a value
		for x, y in self.squares():
			vs = list(map(lambda o: self.var(x, y, o), self.opts()))
			self.add_constraint(vs)

	def add_constraint(self, vs):
		self.constraints.append(vs)

	def add_square(self, x, y, o):
		self.add_constraint([self.var(x, y, o)])

	def print_dimacs(self, fobj):
		vs = set()
		for c in self.constraints:
			for v in c:
				vs.add(abs(v))

		print('p cnf %d %d' % (len(vs), len(self.constraints)), file=fobj)
		for c in self.constraints:
			print('%s 0' % (' '.join(map(str, c)),), file=fobj)

def pipe_to_picosat(sudoku, exe='picosat'):
	ph = subprocess.Popen([exe], stdin=subprocess.PIPE, stdout=subprocess.PIPE, encoding='UTF-8')
	sudoku.print_dimacs(ph.stdin)
	ph.stdin.close()
	ph.wait()
	output = ph.stdout.read()
	solution_vars = []
	for line in output.splitlines():
		if line.startswith('v'):
			vs = line.split()[1:]
			for v in vs:
				v = int(v)
				if v > 0:
					solution_vars.append(v)
	return solution_vars

def print_board(width, height, get_item, fobj=sys.stdout, item_width=2):
	def _print_sep():
		fobj.write("+")
		for x in range(width):
			fobj.write("-" * item_width)
			fobj.write("+")
		fobj.write("\n")

	_print_sep()
	for y in range(height):
		fobj.write("|")
		for x in range(width):
			fobj.write(str(get_item(x, y)).rjust(item_width))
			fobj.write("|")
		fobj.write("\n")
		_print_sep()

s = Sudoku(3)
s.add_base_constraints()
s.add_square(5, 1, 2)
s.add_square(7, 1, 7)
s.add_square(8, 1, 4)
s.add_square(2, 2, 0)
s.add_square(4, 2, 1)
s.add_square(3, 3, 4)
s.add_square(5, 3, 6)
s.add_square(2, 4, 3)
s.add_square(6, 4, 0)
s.add_square(1, 5, 8)
s.add_square(0, 6, 4)
s.add_square(7, 6, 6)
s.add_square(8, 6, 2)
s.add_square(2, 7, 1)
s.add_square(4, 7, 0)
s.add_square(4, 8, 3)
s.add_square(8, 8, 8)

vs = pipe_to_picosat(s)

board = {}
for x, y, o in map(s.decode_var, vs):
	board[x, y] = o

print_board(s.width, s.height, lambda x, y: board[x, y]+1)