pentomino solver

pentomino.py

import string
from copy import deepcopy
from random import shuffle

PENTO = 5
_ = None

# O to Z
piece_labels = string.ascii_uppercase[-12:]

piece_shape = {}
piece_shape["O"] = [["@", "@", "@", "@", "@"]]

piece_shape["P"] = [["@", "@", "@"],
                    ["@", "@"]]

piece_shape["Q"] = [["@", "@", "@", "@"],
                    ["@"]]

piece_shape["R"] = [["@", "@"],
                    [  _, "@", "@"],
                    [  _, "@"]]

piece_shape["S"] = [[  _, "@", "@", "@"],
                    ["@", "@"]]

piece_shape["T"] = [["@"],
                    ["@", "@", "@"],
                    ["@"]]

piece_shape["U"] = [["@", "@", "@"],
                    ["@",   _, "@"]]

piece_shape["V"] = [["@", "@", "@"],
                    ["@"],
                    ["@"]]

piece_shape["W"] = [[  _, "@", "@"],
                    ["@", "@"],
                    ["@"]]

piece_shape["X"] = [[  _, "@"],
                    ["@", "@", "@"],
                    [  _, "@"]]

piece_shape["Y"] = [["@", "@", "@", "@"],
                    [  _, "@"]]

piece_shape["Z"] = [["@", "@"],
                    [  _, "@"],
                    [  _, "@", "@"]]

class Piece:
    def __init__(self, label):
        self.label = label
        shape = piece_shape[label]
        minos = []
        for y, row in enumerate(shape):
            for x, blk in enumerate(row):
                if blk:
                    minos.append((y, x))
        self.minos = minos

    def __eq__(self, other):
        return self.label == other.label

    def __str__(self):
        return self.label

    def possible_positions(self, location, board):
        yspan = len(board)
        xspan = len(board[0])

        # randomize piece direction
        turn_factors = [i for i in range(0, 8)]
        shuffle(turn_factors)

        done = set()
        for factor in turn_factors:
            for center_cell in range(0, PENTO):
                position = MinoPosition(self.minos)
                position.turn_flip(factor)
                position.move(center_cell, location)

                # for deduplicating symmetry
                key = tuple(position.minos)
                if key in done:
                    continue
                done.add(key)

                if position.is_inside(yspan, xspan) and position.can_fit(board):
                    yield position


class MinoPosition:
    def __init__(self, minos):
        self.minos = minos

    def turn_flip(self, factor):
        if factor == 0:
            pass
        elif factor == 1:
            self.minos = [(x, -y) for y, x in self.minos]
        elif factor == 2:
            self.minos = [(-y, -x) for y, x in self.minos]
        elif factor == 3:
            self.minos = [(-x, y) for y, x in self.minos]
        elif factor == 4:
            self.minos = [(y, -x) for y, x in self.minos]
        elif factor == 5:
            self.minos = [(-x, -y) for y, x in self.minos]
        elif factor == 6:
            self.minos = [(-y, x) for y, x in self.minos]
        elif factor == 7:
            self.minos = [(x, y) for y, x in self.minos]
        self.minos.sort()

    def move(self, basecell, center):
        y0, x0 = self.minos[basecell]
        self.minos = [(cy - y0 + center[0], cx - x0 + center[1]) for cy, cx in self.minos]

    def is_inside(self, yspan, xspan):
        return all([0 <= m[0] < yspan and 0 <= m[1] < xspan for m in self.minos])

    def can_fit(self, board):
        return all([board[m[0]][m[1]] is None for m in self.minos])

    def __repr__(self):
        return "MinoPosition(" + repr(self.minos) +")"

    def __iter__(self):
        return iter(self.minos)

pieces = [Piece(s) for s in piece_labels]

class NoAnswer(Exception):
    def __init__(self, board):
        # board not a solition
        self.board = board

trial = 0

def backtrack(pointer, remaining_pieces, board):
    global trial
    ny = len(board)
    nx = len(board[0])

    while pointer < nx * ny:
        y = pointer // nx
        x = pointer % nx

        # find the next empty cell
        if board[y][x] is not None:
            pointer += 1
            continue

        # try all pieces and positions
        for piece in remaining_pieces:
            for piece_position in piece.possible_positions([y, x], board):
                board2 = deepcopy(board)
                for my, mx in piece_position:
                    board2[my][mx] = piece.label
                remaining_pieces2 = [p for p in remaining_pieces if p != piece]

                try:
                    return backtrack(pointer + 1, remaining_pieces2, board2)
                except NoAnswer as e:
                    trial += 1

                    # print progress
                    if trial % 1000 == 0:
                        print("#trial:", trial)
                        print_board(e.board)
                        print()

                    continue

        raise NoAnswer(board)

    return board

def print_board(board):
    for row in board:
        print(" ".join([c if c else "." for c in row]))

# should YSPAN >= XSPAN. it makes calc faster.
YSPAN = 8
XSPAN = 8
plain_board = [[None] * XSPAN for _ in range(YSPAN)]

if __name__ == "__main__":
    board = deepcopy(plain_board)
    for y, x in [(3,3), (3,4), (4,3), (4,4)]:
        board[y][x] = "#"
    print_board(board)
    shuffle(pieces)
    answer = backtrack(0, pieces, board)
    print("#trial:", trial)
    print_board(answer)