LI queens in z3

linkedinqueens.py

# Associated newsletter: https://buttondown.com/hillelwayne/archive/solving-linkedin-queens-with-smt/
# See also: https://ryanberger.me/posts/queens/
# And: https://github.com/ryan-berger/queens/blob/master/main.py

from z3 import * # type: ignore
from itertools import combinations, chain, product
solver = Solver()
size = 9

# queens[n] = col of queen on row n
# by construction, not on same row
queens = IntVector('q', size) 
solver.add([And(0 <= i, i < size) for i in queens])

# not on same column
solver.add(Distinct(queens))

# not diagonally adjacent
for i in range(size-1):
    q1, q2 = queens[i], queens[i+1]
    solver.add(Abs(q1 - q2) != 1)

regions = {
        "purple": [(0, 0), (0, 1), (0, 2), (0, 3), (0, 4), (0, 5), (0, 6), (0, 7), (0, 8),
                   (1, 0), (2, 0), (3, 0), (4, 0), (5, 0), (6, 0), (7, 0), (8, 0),
                   (1, 1), (8, 1)],
        "red": [(1, 2), (2, 2), (2, 1), (3, 1), (4, 1), (5, 1), (6, 1), (6, 2), (7, 1), (7, 2), (8, 2), (8, 3),],
        "brown": [(1, 3), (1, 4), (1, 5), (2, 4), (3, 4),],
        "white": [(2, 3), (3, 2), (3, 3), (4, 2), (5, 2), (5, 3), (6, 3),],
        "green": [(7, 3), (7, 4), (7, 5), (8, 4),],
        "yellow": [(4, 3), (4, 4), (4, 5), (5, 4), (6, 4),],
        "orange": [(3, 5), (3, 6), (4, 6), (4, 7), (5, 6), (5, 5), (6, 5),],
        "blue": [(8, 5), (8, 6), (7, 6), (7, 7), (6, 6),],
        "pink": [(1, 6), (1, 7), (1, 8), (2, 5), (2, 6), (2, 7), (2, 8), (3, 7), (3, 8), 
                 (4, 8), (5, 7), (5, 8), (6, 7), (6, 8), (7, 8), (8, 7), (8, 8),]
        }

### 🎶 Our checks, in the middle of our script 🎶
all_squares = set(product(range(size), repeat=2))
def test_i_set_up_problem_right():
    assert all_squares == set(chain.from_iterable(regions.values()))
    
    for r1, r2 in combinations(regions.values(), 2):
        assert not set(r1) & set(r2), set(r1) & set(r2)
    
def render_regions():
    colormap = ["purple",  "red", "brown", "white", "green", "yellow", "orange", "blue", "pink"]
    board = [[0 for _ in range(size)] for _ in range(size)] 
    for (row, col) in all_squares:
        for color, region in regions.items():
            if (row, col) in region:
                board[row][col] = colormap.index(color)+1
        
    for row in board:
        # print(row)
        print("".join(map(str, row)))

render_regions()

for r in regions.values():
    solver.add(Or(
        *[queens[row] == col for (row, col) in r]
        ))


if solver.check() == sat:
    m = solver.model()
    print([(l, m[l]) for l in queens])