bisco/sudoku_z3.py

## sudoku_z3.py
#!/usr/bin/env python3

from z3 import *

#---------------------------#
# const
#---------------------------#
sudoku = []
solver = Solver()


#---------------------------#

def add_rules(n):
    for r in range(n):
        for c in range(n):
            # 各マスには1～nのいずれかの数値が入る
            solver.add(1 <= sudoku[r][c], sudoku[r][c] <= n)
        # 同じ行のマスにはそれぞれ異なる数値が入る
        solver.add(Distinct(sudoku[r]))
    # 同じ列のマスにはそれぞれ異なる数値が入る
    for c in range(n):
        solver.add(Distinct([sudoku[r][c] for r in range(n)]))

    # sqrt(n) x sqrt(n)の同一ブロック内のマスには、それぞれ異なる数値が入る
    block_len = int(pow(n, 0.5))
    for row in range(block_len):
        for col in range(block_len):
            d = []
            for r in range(block_len):
                for c in range(block_len):
                    d.append(sudoku[row*block_len + r][col*block_len + c])
            solver.add(Distinct(d))

def gen_sudoku_map(n):
    """
    the rule of naming each grid
      0  1  2  3  ... n-1
    a a0 a1 a2 a3 ... an-1
    b b0 b1
    c c0
    d
    e
    :
    """
    for row in range(n):
        sudoku.append([])
        r = chr(ord("a") + row)
        for col in range(n):
            key = "{}{}".format(r, col)
            sudoku[row].append(BitVec(key, 8))

def read_sudoku_question(filename):
    """
    file format:
        delimiter: space
        unknown_grid: fill 0
    """
    q = []
    with open(filename) as f:
        row = 0
        for line in f:
            line = line.strip().split()
            q.append([])
            for col, v in enumerate(line):
                v = int(v.strip())
                q[row].append(v)
            row += 1

    return row, q

def subst_var_to_const(q):
    for r, row in enumerate(q):
        for c, v in enumerate(row):
            if v != 0:
                key = "{}{}".format(chr(ord("a") + r), c)
                sudoku[r][c] = substitute(sudoku[r][c], (sudoku[r][c], BitVecVal(v, 8)))


def main():
    if len(sys.argv) != 2:
        print("usage: {} sudoku_question_file".format(sys.argv[0]))
        sys.exit(1)

    n, q = read_sudoku_question(sys.argv[1])
    gen_sudoku_map(n)
    subst_var_to_const(q)
    add_rules(n)
    if solver.check() == sat:
        m = solver.model()
        for r, row in enumerate(q):
            tmp = []
            for c, v in enumerate(row):
                if v == 0:
                    tmp.append(str(m[sudoku[r][c]])) # modelは辞書みたいに使う
                else:
                    tmp.append(str(v))
            print(" ".join(tmp))
    else:
        print("This sudoku question isn't solvable")


if __name__ == "__main__":
    main()
	#!/usr/bin/env python3

	from z3 import *

	#---------------------------#
	# const
	#---------------------------#
	sudoku = []
	solver = Solver()


	#---------------------------#

	def add_rules(n):
	for r in range(n):
	for c in range(n):
	# 各マスには1～nのいずれかの数値が入る
	solver.add(1 <= sudoku[r][c], sudoku[r][c] <= n)
	# 同じ行のマスにはそれぞれ異なる数値が入る
	solver.add(Distinct(sudoku[r]))
	# 同じ列のマスにはそれぞれ異なる数値が入る
	for c in range(n):
	solver.add(Distinct([sudoku[r][c] for r in range(n)]))

	# sqrt(n) x sqrt(n)の同一ブロック内のマスには、それぞれ異なる数値が入る
	block_len = int(pow(n, 0.5))
	for row in range(block_len):
	for col in range(block_len):
	d = []
	for r in range(block_len):
	for c in range(block_len):
	d.append(sudoku[rowblock_len + r][colblock_len + c])
	solver.add(Distinct(d))

	def gen_sudoku_map(n):
	"""
	the rule of naming each grid
	0 1 2 3 ... n-1
	a a0 a1 a2 a3 ... an-1
	b b0 b1
	c c0
	d
	e
	:
	"""
	for row in range(n):
	sudoku.append([])
	r = chr(ord("a") + row)
	for col in range(n):
	key = "{}{}".format(r, col)
	sudoku[row].append(BitVec(key, 8))

	def read_sudoku_question(filename):
	"""
	file format:
	delimiter: space
	unknown_grid: fill 0
	"""
	q = []
	with open(filename) as f:
	row = 0
	for line in f:
	line = line.strip().split()
	q.append([])
	for col, v in enumerate(line):
	v = int(v.strip())
	q[row].append(v)
	row += 1

	return row, q

	def subst_var_to_const(q):
	for r, row in enumerate(q):
	for c, v in enumerate(row):
	if v != 0:
	key = "{}{}".format(chr(ord("a") + r), c)
	sudoku[r][c] = substitute(sudoku[r][c], (sudoku[r][c], BitVecVal(v, 8)))


	def main():
	if len(sys.argv) != 2:
	print("usage: {} sudoku_question_file".format(sys.argv[0]))
	sys.exit(1)

	n, q = read_sudoku_question(sys.argv[1])
	gen_sudoku_map(n)
	subst_var_to_const(q)
	add_rules(n)
	if solver.check() == sat:
	m = solver.model()
	for r, row in enumerate(q):
	tmp = []
	for c, v in enumerate(row):
	if v == 0:
	tmp.append(str(m[sudoku[r][c]])) # modelは辞書みたいに使う
	else:
	tmp.append(str(v))
	print(" ".join(tmp))
	else:
	print("This sudoku question isn't solvable")


	if __name__ == "__main__":
	main()