z3_magic_square.py

from z3 import *

"""
Generate a 3x3 magic square.

Imagine it like so:

-------------
| a | b | c |
-------------
| d | e | f |
-------------
| g | h | i |
-------------
"""
def model_as_square(model):
  output = """ 
-------------
| {} | {} | {} |
-------------
| {} | {} | {} |
-------------
| {} | {} | {} |
-------------""".format(model[a], model[b], model[c],
                        model[d], model[e], model[f],
                        model[g], model[h], model[i])
  return output

s = Solver()

a, b, c, d, e, f, g, h, i = Ints('a b c d e f g h i')
s.add(Distinct(a, b, c, d, e, f, g, h, i)) 
s.add(a > 0, a < 10) 
s.add(b > 0, b < 10) 
s.add(c > 0, c < 10) 
s.add(d > 0, d < 10) 
s.add(e > 0, e < 10) 
s.add(f > 0, f < 10) 
s.add(g > 0, g < 10) 
s.add(h > 0, h < 10) 
s.add(i > 0, i < 10) 
s.add(a + b + c == d + e + f)
s.add(d + e + f == g + h + i)
s.add(g + h + i == a + d + g)
s.add(a + d + g == b + e + h)
s.add(b + e + h == c + f + i)
s.add(c + f + i == a + e + i)
s.add(a + e + i == c + e + g)

if s.check() == sat:
    model = s.model()
print(model_as_square(model))