squares.py

#!/usr/bin/python3

# From github.com/cipherboy/cmsh
# --> `make distclean cms all check install` will give you a local
#     installation.

import cmsh
import itertools

def constraint(puzzle, x, y):
    is_present = puzzle[y][x]

    neighbors = []
    for dx in range(-1, 2, 2):
        if dx + x >= 0 and dx + x < len(puzzle[y]):
            neighbors.append((dx+x, y))
    for dy in range(-1, 2, 2):
        if dy + y >= 0 and dy + y < len(puzzle):
            neighbors.append((x, dy+y))

    is_neighbor_present = False
    for neighbor in neighbors:
        is_neighbor_present |= puzzle[neighbor[1]][neighbor[0]]

    return is_present | is_neighbor_present

def main():
    model = cmsh.Model()
    width  = 9
    height = 4
    pieces = 10

    puzzle = [
        [
            model.var()
            for _ in range(0, width)
        ]
        for _ in range(0, height)
    ]
    vector = model.to_vec(list(itertools.chain.from_iterable(puzzle)))

    valid_puzzle = True
    for y in range(0, height):
        for x in range(0, width):
            model.add_assert(constraint(puzzle, x, y))

    model.add_assert(vector.bit_sum() == pieces)

    if model.solve():
        while model.solve():
            print("==Solution==")
            for row in puzzle:
                output = ''.join(map(lambda x: '-' if not x else 'x', row))
                print(output)
            model.add_assert(model.negate_solution(vector))
    else:
        print("No solution")

if __name__ == "__main__":
    main()