DominoFit z3 solver

solve.py

#!/usr/bin/env python3
import z3
import itertools

SIZE = 7

top = [2, 7, 4, 3, 6, 4, 5]
side = [6, 5, 2, 5, 6, 1, 6]
blacks = [
    (0, 6), (1, 6), (2, 6),
    (1, 3), (2, 3),
    (3, 1), (3, 2),
]

ones = [
    [z3.Bool(f'one_{r}_{c}') for c in range(SIZE)]
    for r in range(SIZE)
]
twos = [
    [z3.Bool(f'two_{r}_{c}') for c in range(SIZE)]
    for r in range(SIZE)
]

def occupancy(r, c):
    return sum([
        ones[r][c],
        twos[r][c],
        ones[r-1][c] if r > 0 else False,
        twos[r][c+1] if c < SIZE - 1 else False,
        (r, c) in blacks,
    ])

def col(arr, c):
    return [arr[r][c] for r in range(SIZE)]

s = z3.Solver()

for i in range(SIZE):
    # row adds to side
    s.add(sum(ones[i]) + 2 * sum(twos[i]) == side[i])
    # col adds to top
    s.add(sum(col(ones, i)) + 2 * sum(col(twos, i)) == top[i])

for r, c in itertools.product(range(SIZE), range(SIZE)):
    # each cell occupied exactly once
    s.add(occupancy(r, c) == 1)

print(s.check())
m = s.model()

print(' ', end='')
for c in range(SIZE):
    print(top[c], end='')
print()

for r in range(SIZE):
    print(f'{side[r]}', end='')
    for c in range(SIZE):
        is_one = m[ones[r][c]]
        is_two = m[twos[r][c]]
        if is_one and is_two:
            print('!', end='')
        elif is_one:
            print('1', end='')
        elif is_two:
            print('2', end='')
        else:
            print('.', end='')
    print()