dancing links

# Solve the exact cover problem. Given a board consisting of a set of "squares", and a set of
# "pieces", each of which covers a given subset of the "squares", choose a subset of the pieces
# such that each square is covered by exactly one piece in the subset.

# The problem is represented by a matrix where the squares are columns and the pieces are rows.

class Problem:  # The main header node of the matrix
    def __init__(self, snames0, cons):
        # snames0 is a list of all square names
        # Cons is a dict mapping each piece to the names of the squares it covers
        self.squares = dict((sname, Header(sname)) for sname in snames0)  # all column headers
        joinrow([self] + self.squares.values())  # connect the control row
        for pname, snames in cons.items():  # add one row for each piece
            nodes = [Node(pname, self.squares[sname]) for sname in snames]
            joinrow(nodes)

    def headers(self):  # all currently active headers
        h = self.right
        while h is not self:
            yield h
            h = h.right

    def solve(self):
        self.solvelist = []
        self.solvestep()
        return self.solvelist  # will be empty if no solution was found

    def solvestep(self):  # recursive function to get one piece of the solution
        if self.right is self: return True  # All squares have been claimed, solution is found
        # choose a square with the fewest number of pieces that can cover it
        square = min(self.headers(), key = lambda h: h.count)
        if square.count == 0: return False  # No way to cover this square; solution is impossible
        for piece in square.pieces():  # Try each of the pieces that can cover this square
            self.solvelist.append(piece.name)  # push this piece onto the solve list
            piece.place()  # place this piece onto the grid and claim the corresponding squares
            if self.solvestep(): # try to solve the reduced problem
                return True   # A solution has been found
            del self.solvelist[-1]  # solution failed; pop this piece off the solve list
            piece.unplace()  # free up the claimed squares
        return False  # No solution was found

class Header:  # A column header, corresponding to a square
    def __init__(self, squarename):
        self.name = squarename
        self.down = self   # start it off in a column by itself
        self.up = self
        self.count = 0   # number of nodes in this column (ie, pieces that can cover this square)

    def pieces(self, rev=False):  # all current nodes in this column
        p = self.up if rev else self.down
        while p is not self:
            yield p
            p = p.up if rev else p.down

    def claim(self):  # call this when a piece claims this square
        for piece in self.pieces(): # remove all pieces that could occupy this square
            piece.remove_row()
        self.right.left = self.left  # remove this column from matrix
        self.left.right = self.right

    def unclaim(self):  # undo a call to claim
        self.right.left = self  # restore this column into matrix
        self.left.right = self
        for piece in self.pieces(rev=True):  # restore pieces that could occupy this square
            piece.restore_row()

class Node:  # A node in the matrix, corresponding to a piece + square
    def __init__(self, piecename, head):
        self.name = piecename
        self.head = head  # the column header for the column this node is in
        self.up = head.up  # add to the bottom of the given column
        self.down = head
        self.restore_vert()

    def squares(self, rev=False):  # all squares this piece covers
        s = self.left if rev else self.right
        while s is not self:
            yield s
            s = s.left if rev else s.right

    def place(self):  # place this piece onto the board and claim the corresponding squares
        self.remove_row()
        for square in self.squares():
            square.head.claim()
        self.head.claim()

    def unplace(self):  # remove this piece from the board (and put it back into the matrix)
        self.head.unclaim()
        for square in self.squares(rev=True):
            square.head.unclaim()
        self.restore_row()

    def remove_vert(self):   # Disconnect this node from the matrix vertically
        self.up.down = self.down
        self.down.up = self.up
        self.head.count -= 1
    def restore_vert(self):  # Connect this node to the matrix vertically
        self.up.down = self
        self.down.up = self
        self.head.count += 1

    def remove_row(self):   # Vertically disconnect all other nodes in this row
        for square in self.squares():
            square.remove_vert()
    def restore_row(self):  # Restore vertical connections of other nodes in this row
        for square in self.squares():
            square.restore_vert()

def joinrow(nodes):
    # Join a list of nodes together horizontally
    for j in range(len(nodes)):
        nodes[j].right = nodes[(j+1)%len(nodes)]
        nodes[j].left = nodes[(j-1)%len(nodes)]

if __name__ == "__main__":
    import random
    # solve the exact cover problem of placing trominoes onto a board with some missing squares
    n, m = 15, 15  # nxn board with m missing squares
    trominoes = [((0,0),(1,0),(2,0)), ((0,0),(0,1),(0,2)), ((0,0),(1,0),(0,1)),
                 ((0,0),(1,0),(1,1)), ((0,0),(0,1),(1,1)), ((1,0),(0,1),(1,1))]

    squares = [(x, y) for x in range(n) for y in range(n)]
    grid = dict((square, ".") for square in squares)  # text representation of board
    for _ in range(m):  # Randomly remove m squares from the board
        square = random.choice(squares)
        squares.remove(square)
        grid[square] = "*"
    def printgrid():
        print "\n".join(" ".join(grid[(x,y)] for x in range(n)) for y in range(n))
        print
    printgrid()

    # Find every possible place you could place a tromino (including some that go off the board)
    pieces = []
    for x0 in range(n):
        for y0 in range(n):
            for tri in trominoes:
                pieces.append(tuple((x0+dx,y0+dy) for dx,dy in tri))
    # Remove any pieces that don't fit
    pieces = [piece for piece in pieces if all(square in squares for square in piece)]
    # A map from each piece to the squares it covers (in this case they're the same thing)
    cons = dict((piece, piece) for piece in pieces)

    # solve the problem
    problem = Problem(squares, cons)
    solution = problem.solve()

    # print the solution
    if solution:
        for c, piece in enumerate(sorted(solution)):
            for square in piece:
                grid[square] = chr(ord("A") + c % 26)

        printgrid()
    else:
        print "No solution found"