gistfile1.py

from copy import copy, deepcopy

class CoverMatrix(object):
    def __init__(self, columns, column_count = None, row_count = None):
        self._cols = deepcopy(columns)
        self.column_count = len(self._cols) if column_count is None else column_count
        self.row_count = len(self._cols[0]) if row_count is None else row_count
        self.deletion_columns = 0L
        self.deletion_rows = 0L

    def __repr__(self):
        return "CoverMatrix({0}, column_count = {1}, row_count = {2})".format(self._cols, self.column_count, self.row_count)

    def __str__(self):
        return '\n'.join(''.join('x' if self[i, j] is None else '1' if self[i, j] else ' ' for j in xrange(self.column_count)) for i in xrange(self.row_count))

    def __getitem__(self, key):
        row, col = key
        if (self.deletion_columns & (1L << col) or
            self.deletion_rows & (1L << row)):
            return None
        return self._cols[col][row]

    def delete_column(self, column):
        self.deletion_columns |= (1L << column)

    def delete_row(self, row):
        self.deletion_rows |= (1L << row)

    def __nonzero__(self):
        return any(self[i, j] for i in xrange(self.row_count) for j in xrange(self.column_count))

    def select(self):
        columns = []
        for col in xrange(self.column_count):
            if self.deletion_columns & (1L << col):
                continue
            c = [self[row, col] for row in xrange(self.row_count)]
            columns.append((col, c))
        i, col = min(columns, key = lambda x: x[1].count(True))
        return i, col

    @staticmethod
    def exact_cover_solve(A):
        if not A:
            print "Solution found!"
            yield []
            return
        import random
        c, c_entries = A.select()
        rows = [n for n, r in enumerate(c_entries) if r]
        random.shuffle(rows)
        for r in rows:
            print "Recursing on action {0} solving constraint {1}...".format(r, c)
            Aprime = copy(A)
            for j in (j for j in xrange(Aprime.column_count) if A[r, j]):
                for i in (i for i in xrange(Aprime.row_count) if A[i, j]):
                    Aprime.delete_row(i)
                Aprime.delete_column(j)
            for x in CoverMatrix.exact_cover_solve(Aprime):
                x.append(r)
                yield x

class ExactCoverProblem(object):
    def __init__(self):
        self.constraints = []
        self.optional_constraints = set()
        self.actions = []

    def add_constraint(self, constraint, optional=False):
        self.constraints.append(constraint)
        if optional:
            self.optional_constraints.add(self.constraints.index(constraint))

    def add_action(self, action, solved_constraints):
        self.actions.append((action, [self.constraints.index(c) for c in solved_constraints]))

    def force_action(self, action):
        constraint = object()
        self.add_constraint(constraint)
        for act in self.actions:
            if act[0] == action:
                act[1].append(self.constraints.index(constraint))
                return

    def solve(self):
        final_actions = copy(self.actions)
        for optional in self.optional_constraints:
            final_actions.append((None, [optional]))
        # Construct the cover matrix
        # Initially filled with all False
        columns = [[False] * len(final_actions) for i in xrange(len(self.constraints))]
        # Then we add on things
        for i, action in enumerate(final_actions):
            for constraint in action[1]:
                columns[constraint][i] = True
        m = CoverMatrix(columns)
        solutions = CoverMatrix.exact_cover_solve(m)
        try:
            rows = solutions.next()
            return [final_actions[n][0] for n in reversed(rows)
                                        if final_actions[n][0]
                                           is not None]
        except StopIteration:
            raise Exception("No solutions")

if __name__ == "__main__":
    import sys
    queens = int(sys.argv[1] if len(sys.argv) > 1 else 8)
    print "Constructing constraints..."
    problem = ExactCoverProblem()
    # First, add constraints: at most one queen in each column and row
    for i in xrange(queens):
        problem.add_constraint(("column", i))
        problem.add_constraint(("row", i))
    # And the diagonals, optionally
    for i in xrange(2*queens - 1):
        problem.add_constraint(("diagonal_right", i), optional=True)
        problem.add_constraint(("diagonal_left", i), optional=True)
    # Add the actions
    print "Constructing actions..."
    for x in xrange(queens):
        for y in xrange(queens):
            if x == y or x == queens - 1 - y:
                continue # Avoid both major diagonals
            problem.add_action((x, y), [("column", y),
                                        ("row", x),
                                        ("diagonal_right", queens - 1 - y + x),
                                        ("diagonal_left", x + y)])
    problem.force_action((0, 3))
    print "Running Algorithm X..."
    solution = problem.solve()
    for y in xrange(queens):
        print ''.join('Q' if (x, y) in solution else '.' for x in xrange(queens))