sundries

 1import z3
 2import sys
 3
 4if len(sys.argv) != 2:
 5    print("usage: magic_square <SIZE>")
 6    exit(1)
 7
 8size = int(sys.argv[1])
 9if size < 3:
10    print("size must be >= 3")
11    exit(1)
12
13s = z3.Solver()
14
15vars = []
16for i in range(0, size):
17    row = []
18    for j in range(0, size):
19        v = z3.Int(f"{i}{j}")
20        s.add(v >= 1)
21        row.append(v)
22    vars.append(row)
23
24for i in range(0, size):
25    for j in range(0, size):
26        for k in range(0, size):
27            for l in range(0, size):
28                if i == k and j == l:
29                    continue
30                s.add(vars[i][j] != vars[k][l])
31
32# rows
33exps = []
34for i in range(0, size):
35    exp = vars[i][0]
36    for j in range(1, size):
37        exp = exp + vars[i][j]
38    exps.append(exp)
39
40# cols
41for j in range(0, size):
42    exp = vars[0][j]
43    for i in range(1, size):
44        exp = exp + vars[i][j]
45    exps.append(exp)
46
47# diags
48one = vars[0][0]
49two = vars[0][2]
50for i in range(1, size):
51    one = one + vars[i][i]
52    two = two + vars[i][size - i - 1]
53
54exps.append(one)
55exps.append(two)
56
57for i in range(0, len(exps) - 1):
58    s.add(exps[i] == exps[i + 1])
59
60if s.check():
61    model = s.model()
62    for i in range(0, size):
63        for j in range(0, size):
64            print("{}\t".format(model[vars[i][j]]), end="")
65        print()
66else:
67    print("Model wasn't satisfyable")