lub2022-practical/ex3/trees_and_tents.py

# coding: utf-8
import os, sys
from z3 import *

# get the playground information
if len(sys.argv) != 2:
    print("Usage: python3 trees_and_tents.py <forest-file>")

EMPTY = "?"
TREE = 2

if len(sys.argv) == 2:
    fname = sys.argv[1]

with open(fname) as f:
    playground = f.read()
    cols = [col for col in playground.strip().split("\n")[0][1:]]
    rows = [row[0] for row in playground.strip().split("\n")[1:]]
playground = [[EMPTY if x == EMPTY else x for x in row[1:]]
        for row in playground.strip().split("\n")[1:]]

# get the playground size
len_rows = len(rows)
assert(len_rows != 0)
len_cols = len(cols)
assert(len_cols != 0)

def print_result(model):
    print("  " + ''.join(cols) + " "*5 + ''.join(cols))
    for i in range(len_rows):
        print(rows[i] + " " + ''.join(playground[i]) + " >>> ", end='')
        for j in range(len_cols):
            cell = model[cells[i][j]].as_long()
            if cell != TREE:
                print(cell, end='')
            elif cell == TREE:
                print("T", end='')
        print()

def print_forest():
    print("  " + ''.join(cols))
    for i in range(len_rows):
        print(rows[i] + " ", end='')
        for j in range(len_cols):
            cell = playground[i][j]
            if cell != TREE:
                print(cell, end='')
            elif cell == TREE:
                print("T", end='')
        print()

################################# Trees and Tents ##################################
# create the solver
solver = Solver()
cells = [[None for j in range(len_cols)] for i in range(len_rows)]

def get_possible_tents_in_col(col):
    possible_tent_positions = []
    for row in range(len_rows):
        if playground[row][col] != "T":
           possible_tent_positions.append(cells[row][col])
    return possible_tent_positions

def get_possible_tents_in_row(row):
    possible_tent_positions = []
    for col in range(len_cols):
        if playground[row][col] != "T":
           possible_tent_positions.append(cells[row][col])
    return possible_tent_positions

def get_neighbours(i, j):
    cs = []
    for p in range(max(i-1, 0), min(i+2, len_rows)):
        for q in range(max(j-1, 0), min(j+2, len_cols)):
            if p == i and q == j: continue
            cs.append(cells[p][q])
    return cs

def get_tree_neighbours(i, j):
    tree_pos = []
    for p in [max(i-1, 0), min(i+1, len_rows - 1)]:
        if p == i: continue
        tree_pos.append(cells[p][j])
    for q in [max(j-1, 0), min(j+1, len_cols - 1)]:
        if q == j: continue
        tree_pos.append(cells[i][q])
    return tree_pos

# TODO
# Cell entries need to be represented by a Z3 variable
for i in range(len_rows):
    for j in range(len_cols):
    #   your code goes here...
# TODO
# Bound/restrict the values of the cells according to the input
        if playground[i][j] == EMPTY:
        #   your code goes here...
        else:
        #   your code goes here...

# TODO
# The sums of tents per row/column must match the amounts given in the input
for j in range(len_cols):
#   your code goes here...
for i in range(len_rows):
#   your code goes here...

for i in range(len_rows):
    for j in range(len_cols):
        if playground[i][j] == "T": continue
        tree_constraint = Bool("False")
        # TODO
        # A tent needs to be next to a tree
        # Incrementally build up the tree_constraint ...
        for possible_tree in get_tree_neighbours(i, j):
        #   your code goes here...
        # ... and add it to the solver
        # TODO
        # A tent must not be next to another tent
        for neighbour in get_neighbours(i, j):
        #   your code goes here...

res = solver.check()
if res == unsat:
    print("UNSAT")
    print_forest()
else:
    m = solver.model()
    print_result(m)