tents_and_trees_solver.py

# pip install z3-solver
from z3 import Int, Solver, And, sat, If, Sum, Or
from typing import List, Tuple

width = 8  # Number of columns
height = 8  # Number of rows

# Matrix of integers (0 if grass, 1 if tent), using integer so we can use Sum to count them
X = [[Int(f"x_{i+1}_{j+1}") for j in range(width)] for i in range(height)]

col_vals = (2, 1, 2, 1, 1, 2, 1, 2)
row_vals = (3, 1, 2, 1, 1, 1, 2, 1)

assert len(col_vals) == width and len(row_vals) == height

# 1 if there's a tree there
trees = (
    (0, 1, 0, 0, 0, 0, 0, 0),
    (1, 0, 0, 0, 0, 1, 1, 0),
    (0, 0, 0, 0, 0, 0, 0, 0),
    (0, 1, 0, 1, 0, 0, 0, 0),
    (0, 0, 0, 0, 0, 1, 1, 0),
    (0, 0, 1, 0, 0, 0, 0, 0),
    (0, 0, 0, 0, 0, 1, 0, 0),
    (1, 0, 0, 0, 0, 0, 1, 0),
)

assert len(trees) == height and all(len(line) == width for line in trees)

# Every integer can only be 0 (grass) or 1 (tent)
bool_c = [Or(X[i][j] == 0, X[i][j] == 1) for i in range(height) for j in range(width)]

# Constraint on number of tents in each row
rows_c = [Sum(X[i]) == row_vals[i] for i in range(height)]

# Constraint on number of tents in each column
cols_c = [Sum([X[i][j] for i in range(height)]) == col_vals[j] for j in range(width)]


# Every house must be adjacent to exactly one tree and no other houses


def square_adjacents(row: int, col: int) -> List[Tuple[int, int]]:
    # Returns square adjancent coorinates inside the square
    all_adjacents = [(row + 1, col), (row - 1, col), (row, col + 1), (row, col - 1)]
    return [
        (adj_row, adj_col)
        for adj_row, adj_col in all_adjacents
        if 0 <= adj_row < height and 0 <= adj_col < width
    ]


def diagonal_adjacents(row: int, col: int) -> List[Tuple[int, int]]:
    # Returns square and diagonal adjancent coorinates inside the square
    diag_adjacents = [
        (row + 1, col + 1),
        (row + 1, col - 1),
        (row - 1, col + 1),
        (row - 1, col - 1),
    ]
    return [
        (adj_row, adj_col)
        for adj_row, adj_col in diag_adjacents
        if 0 <= adj_row < height and 0 <= adj_col < width
    ] + square_adjacents(row, col)


adj_c = []
for row in range(height):
    for col in range(width):
        square_adj = square_adjacents(row, col)
        diag_adj = diagonal_adjacents(row, col)

        if trees[row][col]:
            # If trees[row][col], X[row][col] must be grass
            adj_c.append(X[row][col] == 0)
            # At least one adjacent cell must have a tent
            adj_c.append(
                Or([X[adj_row][adj_col] == 1 for adj_row, adj_col in square_adj])
            )
        else:
            # If X[row][col] has a tent, at least one adjacent must be a tree (static constraint)
            # and no adjacent cell must be a tent (dynamic constraint)
            adj_c.append(
                If(
                    X[row][col] == 1,
                    And(
                        any(trees[adj_row][adj_col] for adj_row, adj_col in square_adj),
                        And(
                            [X[adj_row][adj_col] == 0 for adj_row, adj_col in diag_adj]
                        ),
                    ),
                    True,
                )
            )

s = Solver()
s.add(bool_c + rows_c + cols_c + adj_c)
if s.check() == sat:
    m = s.model()
    solution = [[m.evaluate(X[i][j]) for j in range(width)] for i in range(height)]
    for row, solution_row in enumerate(solution):
        grasses = ["🟩", "🌲"]
        print(
            "".join(
                "⛺" if v == 1 else grasses[trees[row][col]]
                for col, v in enumerate(solution_row)
            )
        )
else:
    print("failed to solve")