lub2022-practical/ex2/square.py

# coding: utf-8
import os, sys
from z3 import *

# get the playground information
if len(sys.argv) != 2:
    print("Usage: python3 square.py <level-file>")
    sys.exit(0)

with open(sys.argv[1]) as f:
    playground = f.read()
rows = playground.strip().split("\n")
playground = [[None if x == "_" else int(x) for x in r.split()] for r in rows]

# get the playground size
size_y = len(playground)
assert(size_y != 0)
size_x = len(playground[0])
assert(size_x != 0)
assert(size_x == size_y)

#################################### Square ####################################

# create the solver
solver = Solver()

# todo: create an integer variable for each playground cell
# hint: use something like the coordinates as part of the variable name
numbers = [[None for _j in range(size_x)] for _j in range(size_y)]
for i in range(size_y):
    for j in range(size_x):
        pass # ... replace this with your code ...

# todo: assign each known number the corresponding value from playground
for i in range(size_y):
    for j in range(size_x):
        pass # ... replace this with your code ...

# todo: declare a variable for the sum of all columns, rows and diagonals
# ... your code goes here ...

# todo: enforce that each column sums up to the declared variable
for j in range(size_x):
    pass # ... replace this with your code ...

# todo: enforce that each row sums up to the declared variable
for i in range(size_y):
    pass # ... replace this with your code ...

# todo: enforce that both diagonals sum up to the declared variable

# call the solver and check satisfiability
res = solver.check()
if res != sat:
    print("unsat")
    sys.exit(1)

# print the model
m = solver.model()
for i in range(size_y):
    results = []
    for j in range(size_x):
        num = numbers[i][j]
        results.append("_" if num is None else m[num].as_long())
    print(("%4s" * len(results)) % tuple(results))

################################################################################