Use back tracking to solve peg solitaire (English board)

peg_solitaire_solver.py

import sys

class PegSolitaireBoard:
    """Represents a peg solitaire board

    """
    def __init__(self):
        """Constructor

        Returns:
            board (PegSolitaireBoard):
        """

        # 1 means occupied, 0 means free, 2 means out of bounds
        self.table = [
            [2, 2, 1, 1, 1, 2, 2],
            [2, 2, 1, 1, 1, 2, 2],
            [1, 1, 1, 1, 1, 1, 1],
            [1, 1, 1, 0, 1, 1, 1],
            [1, 1, 1, 1, 1, 1, 1],
            [2, 2, 1, 1, 1, 2, 2],
            [2, 2, 1, 1, 1, 2, 2]
        ]

    def reset(self):
        """Rests the board to its starting state

        """
        self.__init__()

    def move(self, row, col, dir):
        """Attempt to move a peg at a position in a direction.

        This method will check to see if there is a peg at the specified position, and if
        the move is legal. The move will only be made if it's legal.

        Args:
            row (int): Row of peg
            col (int): Column of peg
            dir (int): Direction to move.
                Possible choices are 0 (left), 1 (right), 2 (up), and 3 (down)

        Returns:
            success (bool): True if it's there is a peg and the move is made, False if
                the move is illegal
        """
        if self.table[row][col] != 1:
            return False

        if dir == 0 and col - 2 >= 0\
                    and self.table[row][col - 1] == 1\
                    and self.table[row][col - 2] == 0:
            self.table[row][col] = self.table[row][col - 1] = 0
            self.table[row][col - 2] = 1

            return True

        if dir == 1 and col + 2 < len(self.table[0]) \
                    and self.table[row][col + 1] == 1 \
                    and self.table[row][col + 2] == 0:
            self.table[row][col] = self.table[row][col + 1] = 0
            self.table[row][col + 2] = 1

            return True

        if dir == 2 and row - 2 >= 0 \
                    and self.table[row - 1][col] == 1 \
                    and self.table[row - 2][col] == 0:
            self.table[row][col] = self.table[row - 1][col] = 0
            self.table[row - 2][col] = 1

            return True

        if dir == 3 and row + 2 < len(self.table) \
                    and self.table[row + 1][col] == 1 \
                    and self.table[row + 2][col] == 0:
            self.table[row][col] = self.table[row + 1][col] = 0
            self.table[row + 2][col] = 1

            return True

        return False

    def unmove(self, row, col, dir):
        """Undo the move of a peg in a direction.

        Method has no sanity checking, is only safe to call if the last call to move
        returns True.

        Args:
            row (int): Row of peg
            col (int): Column of peg
            dir (int): Direction to move.
        """
        if dir == 0:
            self.table[row][col] = self.table[row][col - 1] = 1
            self.table[row][col - 2] = 0
            return

        if dir == 1:
            self.table[row][col] = self.table[row][col + 1] = 1
            self.table[row][col + 2] = 0
            return

        if dir == 2:
            self.table[row][col] = self.table[row - 1][col] = 1
            self.table[row - 2][col] = 0

        if dir == 3:
            self.table[row][col] = self.table[row + 1][col] = 1
            self.table[row + 2][col] = 0

    def pretty_print(self):
        """Print the board nicely

        """
        output_map = {
            2: ' x ',
            0: '   ',
            1: ' i '
        }

        rows = len(self.table)
        cols = len(self.table[0])

        sys.stdout.write('  ')
        for col_ind in xrange(cols):
            sys.stdout.write(' {0} '.format(str(col_ind)))

        sys.stdout.write('\n')

        for row_ind in xrange(rows):
            for col_ind in xrange(cols):
                if col_ind == 0:
                    sys.stdout.write('{0} '.format(str(row_ind)))

                sys.stdout.write('{0}'.format(output_map[self.table[row_ind][col_ind]]))

                if col_ind == cols - 1:
                    sys.stdout.write(' \n')


class PegSolitaireSolver:
    def __init__(self, board, relax=False):
        """Constructor

        Args:
            board (PegSolitaireBoard): The peg solitaire board
            relax (bool): Set to true if position of last peg doesn't matter, default
                is False

        Returns:
            solver (PegSolitaireSolver):

        """
        self.board = board
        self.relax = relax

        num_pegs = 0
        for row in self.board.table:
            for elem in row:
                if elem == 1:
                    num_pegs += 1

        # Max number of moves that the board allows
        self.max_moves = num_pegs - 1

        # Coordinates of the center of the board
        self.center_row = len(self.board.table) / 2
        self.center_column = len(self.board.table[0]) / 2

        self.solution = []

    def _back_track(self, move):
        """Do back tracking to see if the board has a solution

        Args:
            move (int): The move number

        Returns:
            success (bool): True if a solution exists

        """
        # Base case
        if move == self.max_moves:
            # Since each move removes 1 peg, only check if the center has a peg
            if self.relax or self.board.table[self.center_row][self.center_column] == 1:
                return True
            else:
                return False

        # Iterate through the board positions
        for r in xrange(len(self.board.table)):
            for c in xrange(len(self.board.table[0])):
                # Attempt to make a move in each direction
                for d in xrange(4):
                    # Check if a legal move is made
                    if self.board.move(r, c, d):
                        # Check if branch leads to solution
                        if self._back_track(move + 1):
                            # Save the move if a solution is found
                            self.solution.append((r, c, d))
                            return True

                        # Undo the move if branch has no solution
                        self.board.unmove(r, c, d)

        return False

    def solve(self):
        """Find the series of moves that solves the solitaire board

        Returns:
            moves (list): A list of moves that will solve the board.
                Each element is of form (row, column, direction), and are valid arguments
                to the move method of a PegSolitaireBoard object

        """
        # Reset the board and the solution list
        self.solution = []
        self.board.reset()

        # Start back tracking
        self._back_track(0)

        # Solution needs to be reversed
        self.solution.reverse()

        return self.solution

    def print_solution(self, interactive=False):
        """Apply the solution to the starting board, and print the state of the board
        after each step.

        """
        self.board.reset()

        if interactive:
            print chr(27) + "[2J"

        print '\nStarting board\n'
        self.board.pretty_print()

        dir_map = {
            0: 'left',
            1: 'right',
            2: 'up',
            3: 'down'
        }
        for step in self.solution:
            r, c, d = step

            if interactive:
                raw_input('\nPress enter for next move')
                print chr(27) + "[2J"

            print '\nMoved row {0} column {1} {2}\n'.format(r, c, dir_map[d])

            self.board.move(*step)
            self.board.pretty_print()


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '-i', '--interactive',
        help='Print solution in interactive mode',
        action='store_true'
    )
    parser.add_argument(
        '-r', '--relax',
        help='Position of last peg doesn\'t matter',
        action='store_true'
    )
    args = parser.parse_args()

    bc = PegSolitaireSolver(PegSolitaireBoard(), relax=args.relax)
    bc.solve()

    bc.print_solution(args.interactive)