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(queens):
    revert_color = "\x1b[39;49m"
    # Using U+274C CROSS MARK as a missing tile indicator. It's a Wide character so it should take two columns in a terminal
    missing = revert_color + "\u274c\ufe0f"  
    empty = "  "
    # The U+265B BLACK CHESS QUEEN is a Narrow character; it's also not an emoji, and does not have an emoji presentation. So
    # we add a space after it so it will take two columns in a terminal. We use BLACK CHESS QUEEN instead of WHITE CHESS QUEEN
    # because it usually has less spacings inside its glyph in the fonts, so it will have more foreground color in it when rendered
    queen = "\u265b \x1b[39m"
    # Colours in terminals are stupidly convoluted for historical reasons. The queens are rendered black on brown, white, yellow,
    # orange, and blue tiles, and white on purple, red, green, and blue tiles, for better visual contrast and ease of observation
    colormap = {
        "purple": "\x1b[97;48:5:53m",
        "red": "\x1b[97;41m",
        "brown": "\x1b[30;43m",
        "white": "\x1b[30;107m",
        "green": "\x1b[97;42m",
        "yellow": "\x1b[30;103m",
        "orange": "\x1b[30;48:5:202m",
        "blue": "\x1b[97;44m",
        "pink": "\x1b[30;48:5:210m",
    }
    board = [[missing 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:
                tile = queen if (row, col) in queens else empty
                board[row][col] = colormap[color] + tile

    for row in board:
        print("".join(row) + revert_color)

render_regions([])
print()

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


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